高性能PyTorch是如何炼成的？整理的10条脱坑指南

将 RGB 图像保持在每个通道深度 8 位。可以轻松地在 GPU 上将图像转换为浮点形式或者标准化。 在数据集中用 uint8 或 uint16 数据类型代替 long。 class MySegmentationDataset(Dataset):   ...   def __getitem__(self, index):     image = cv2.imread(self.images[index])     target = cv2.imread(self.masks[index])      # No data normalization and type casting here     return torch.from_numpy(image).permute(2,0,1).contiguous(),            torch.from_numpy(target).permute(2,0,1).contiguous()  class Normalize(nn.Module):     # https://github.com/BloodAxe/pytorch-toolbelt/blob/develop/pytorch_toolbelt/modules/normalize.py     def __init__(self, mean, std):         super().__init__()         self.register_buffer("mean", torch.tensor(mean).float().reshape(1, len(mean), 1, 1).contiguous())         self.register_buffer("std", torch.tensor(std).float().reshape(1, len(std), 1, 1).reciprocal().contiguous())      def forward(self, input: torch.Tensor) -> torch.Tensor:         return (input.to(self.mean.type) - self.mean) * self.std  class MySegmentationModel(nn.Module):   def __init__(self):     self.normalize = Normalize([0.221 * 255], [0.242 * 255])     self.loss = nn.CrossEntropyLoss()    def forward(self, image, target):     image = self.normalize(image)     output = self.backbone(image)      if target is not None:       loss = self.loss(output, target.long())       return loss      return output 

通过这样做，会大大减少 RAM 的需求。对于上面的示例。用于高效存储数据表示的内存使用量将为每批 33Mb，而之前是 167Mb，减少为原来的五分之一。当然，这需要模型中添加额外的步骤来标准化数据或将数据转换为合适的数据类型。但是，张量越小，CPU 到 GPU 的传输就越快。

DataLoader 的工作程序的数量应该谨慎选择。你应该查看你的 CPU 和 IO 系统有多快，你有多少内存，GPU 处理数据有多快。

多 GPU 训练 & 推理神经网络模型变得越来越大。今天，使用多个 GPU 来增加训练时间已成为一种趋势。幸运的是，它经常会提升模型性能来达到更大的批处理量。PyTorch 仅用几行代码就可以拥有运行多 GPU 的所有功能。但是，乍一看，有些注意事项并不明显。

model = nn.DataParallel(model) # Runs model on all available GPUs 

运行多 GPU 最简单的方法就是将模型封装在 nn.DataParallel 类中。除非你要训练图像分割模型(或任何生成大型张量作为输出的其他模型)，否则大多数情况下效果不错。在正向推导结束时，nn.DataParallel 将收集主 GPU 上所有的 GPU 输出，来通过输出反向运行，并完成梯度更新。

于是，现在就有两个问题：

GPU 负载不平衡; 在主 GPU 上聚合需要额外的视频内存

首先，只有主 GPU 能进行损耗计算、反向推导和渐变步骤，其他 GPU 则会在 60 摄氏度以下冷却，等待下一组数据。

其次，在主 GPU 上聚合所有输出所需的额外内存通常会促使你减小批处理的大小。nn.DataParallel 将批处理均匀地分配到多个 GPU。假设你有 4 个 GPU，批处理总大小为 32;然后，每个 GPU 将获得包含 8 个样本的块。但问题是，尽管所有的主 GPU 都可以轻松地将这些批处理放入对应的 VRAM 中，但主 GPU 必须分配额外的空间来容纳 32 个批处理大小，以用于其他卡的输出。

对于这种不均衡的 GPU 使用率，有两种解决方案：

（编辑：晋中站长网）

【声明】本站内容均来自网络，其相关言论仅代表作者个人观点，不代表本站立场。若无意侵犯到您的权利，请及时与联系站长删除相关内容!