PyTorch 网络单集运行时间计算 (PyCharm)

要计算网络单集的运行时间，可以使用 PyTorch 中的 'torch.utils.benchmark' 模块。以下是一个简单的示例：

import torch
import timeit

# 定义网络
class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = torch.nn.Linear(784, 256)
        self.fc2 = torch.nn.Linear(256, 10)

    def forward(self, x):
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = torch.relu(x)
        x = self.fc2(x)
        output = torch.log_softmax(x, dim=1)
        return output

# 创建网络实例
net = Net()

# 创建输入数据
input_data = torch.randn(64, 1, 28, 28)

# 计算单集运行时间
timer = timeit.Timer(lambda: net(input_data))
print(timer.timeit(number=1))

在这个示例中，我们首先定义了一个网络类 'Net'，然后创建了一个网络实例 'net' 和输入数据 'input_data'，最后使用 'timeit.Timer' 计算了网络单集的运行时间。需要注意的是，'timeit.Timer' 的 'number' 参数表示运行的次数，这里设置为 1，即只运行一次。