当网络需要共享（常量）`Tensor`时使用`DataParallel`

Question

我想使用DataParallel 将我的计算分布在批处理维度上的多个 GPU 上。我的网络在内部需要一个Tensor（我们称之为A），它是恒定的，不会通过优化而改变。似乎DataParallel 不会自动将此Tensor 复制到所有相关的GPU，因此网络会抱怨它看到的输入数据块x 驻留在与A 不同的GPU 上。

有没有办法DataParallel 可以自动处理这种情况？或者，有没有办法将Tensor 复制到所有 GPU？还是应该只为每个 GPU 保留一个 Tensor，然后根据 forward 看到的块所在的位置手动确定要使用哪个副本？

Answer 1

您应该将您的张量包装在 torch.nn.Parameter 中，并在创建过程中设置 requires_grad=False。

torch.nn.Parameter 并不意味着张量必须是可训练的。

这仅仅意味着它是模型的一部分，如果需要（例如多个 GPU）应该被转移。

如果不是这样，torch 就无法知道__init__ 中的哪个张量是模型的一部分（您可以对张量进行一些操作并添加到self 以完成某些工作)。

我认为不需要其他函数来做到这一点，尽管名称可能有点混乱。