如何编写一个快速的 PyTorch 训练循环？答案

【问题标题】：How to write a fast PyTorch training loop?如何编写一个快速的 PyTorch 训练循环？
【发布时间】：2021-08-15 02:01:45
【问题描述】：

在将机器学习算法从 Keras 移植到 PyTorch 后，我遇到了性能下降的问题。

Keras 中的以下构造：

model.compile(loss="binary_crossentropy", optimizer=SGD(learning_rate=0.01)) 
.
.
.
model.fit(states, actions)

比 Pytorch 中的以下详细版本快约 15 倍：

def train_network(model, optimizer, train_loader,
                  num_epochs=16, pbar_update_interval=200, print_logs=True):

    criterion = nn.BCELoss()
    model.train()

    for i in range(num_epochs):
        for k, batch_data in enumerate(train_loader):
            optimizer.zero_grad()

            batch_x = batch_data[:, :-1]
            batch_y = batch_data[:, -1]

            y_pred = model(batch_x)
            loss = criterion(y_pred, batch_y.unsqueeze(1))
            loss.backward() 
            optimizer.step() 

optimizer = torch.optim.SGD(model.parameters(), lr=LEARNING_RATE)
train_data = torch.from_numpy(np.column_stack((states, actions)))
train_data = train_data.to(torch.float)
train_loader = torch.utils.data.DataLoader(train_data, shuffle=True, batch_size=32)
train_network(model, optimizer, train_loader)

我是不是写错了 PyTorch 代码？

PyTorch 代码中的变量model 是nn.Module，而states 和actions 是NumPy 数组。代码在 4 个 CPU 内核上运行。

【问题讨论】：

尝试测量拆分batch_x和y的时间。这些线产生了另外 2 个张量，所以如果它是一个非常小的模型，那么这可能就是原因。
您需要将实际模型放在这里。取决于您要测试的模型类型。 Pytorch 中的小模型有点慢。并且有些层比 TF 慢，而有些层则更快。

标签： python machine-learning keras optimization pytorch

【解决方案1】：

keras.model.fit 似乎有默认参数值epochs=1。所以这正好解释了性能因素。

【讨论】：