Flops及训练时间的测试代码

以下是一个用于测试模型 FLOPs 和训练时间的代码示例：

import torch
import time

# 定义模型
class MyModel(torch.nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.conv1 = torch.nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1)
        self.conv2 = torch.nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1)
        self.pool = torch.nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = torch.nn.Linear(128 * 8 * 8, 1024)
        self.fc2 = torch.nn.Linear(1024, 10)

    def forward(self, x):
        x = torch.nn.functional.relu(self.conv1(x))
        x = self.pool(x)
        x = torch.nn.functional.relu(self.conv2(x))
        x = self.pool(x)
        x = x.view(-1, 128 * 8 * 8)
        x = torch.nn.functional.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# 测试模型 FLOPs
model = MyModel()
inputs = torch.rand(1, 3, 32, 32)
flops, params = torch.profiler.profile(model, inputs=(inputs,))
print(f"FLOPs: {flops / 1e9:.2f}G")
print(f"Params: {params / 1e6:.2f}M")

# 测试模型训练时间
optimizer = torch.optim.Adam(model.parameters())
loss_fn = torch.nn.CrossEntropyLoss()
train_data = torch.utils.data.DataLoader(torch.randn(1000, 3, 32, 32), batch_size=100)
start_time = time.time()
for epoch in range(10):
    for x in train_data:
        optimizer.zero_grad()
        pred = model(x)
        loss = loss_fn(pred, torch.randint(0, 10, size=(100,)))
        loss.backward()
        optimizer.step()
print(f"Training time: {time.time() - start_time:.2f}s")

这个示例中，我们定义了一个简单的卷积神经网络模型 `MyModel`，并使用 `torch.profiler` 模块计算了模型的 FLOPs 和参数数量。然后我们使用随机数据进行了简单的模型训练，并计算了训练时间。注意这里的训练时间仅供参考，实际训练时间取决于许多因素，例如 GPU/CPU 的性能、数据集的大小等。