Gluon 代码在 GPU 上的速度并不比在 CPU 上快多少

Question

我们正在为三元组训练推荐系统的网络。 fit方法的核心代码如下：

for e in range(epochs):
    start = time.time()

    cumulative_loss = 0

    for i, batch in enumerate(train_iterator):
        # Forward + backward.
        with autograd.record():
            output = self.model(batch.data[0])
            loss = loss_fn(output, batch.label[0])

        # Calculate gradients
        loss.backward()
        # Update parameters of the network.
        trainer_fn.step(batch_size)
        # Calculate training metrics. Sum losses of every batch.
        cumulative_loss += nd.mean(loss).asscalar()
    train_iterator.reset()

train_iterator 是一个自定义迭代器类，它继承自 mx.io.DataIter，并在适当的上下文中返回数据（三元组），如下所示：

        data = [mx.nd.array(data[:, :-1], self.ctx, dtype=np.int)]
        labels = [mx.nd.array(data[:, -1], self.ctx)]
        return mx.io.DataBatch(data, labels)

self.model.initialize(ctx=mx.gpu(0)) 在运行fit 方法之前也被调用。 loss_fn = gluon.loss.L1Loss().

问题在于nvidia-smi 报告进程已正确分配到 GPU。但是，在 GPU 中运行 fit 并不比在 CPU 中运行快多少。此外，将batch_size 从 50000 增加到 500000 会使每批次的时间增加 10 倍（考虑到 GPU 并行化，这是我没想到的）。

具体来说，对于 50k 批次： * output = self.model(batch.data[0]) 在 GPU 上耗时 0.03 秒，在 CPU 上耗时 0.08 秒。 * loss.backward() 耗时 0.11 秒，CPU 耗时 0.39。

两者都使用nd.waitall() 进行评估，以避免异步调用导致错误测量。

此外，在普通 MXNet 上运行的非常相似的代码在相应部分上花费的时间不到 0.03 秒，这导致一个完整的 epoch 需要从 MXNet 略高于 1 分钟的时间到 Gluon 的 15 分钟。

对这里可能发生的事情有什么想法吗？

提前致谢！

Answer 1

问题出在下面一行：

cumulative_loss += nd.mean(loss).asscalar()

当您调用asscalar() 时，MXNet 必须隐式执行同步调用以将结果从 GPU 复制到 CPU：这与调用 nd.waitall() 基本相同。由于您为每次迭代都执行此操作，因此每次迭代都会进行同步，从而显着降低您的挂钟时间。

您可以做的是在 GPU 中保留和更新您的 cumulative_loss 并仅在您实际需要显示它时将其复制到 CPU - 它可以是每 N 次迭代或在 epoch 实际完成后，具体取决于多长时间每次迭代都需要完成。