收藏丨高性能PyTorch是如何炼成的？过来人吐血整理的避坑指南 慧维智能 慧维智能

选自towardsdatascience 作者：Eugene Khvedchenya

机器之心编译，参与：小舟、蛋酱、魔王。

class RAMDataset(Dataset):
   def __init__(image_fnames, targets):
    self.targets = targets
    self.images = []
    for fname in tqdm(image_fnames, desc="Loading files in RAM"):
      with open(fname, "rb") as f:
        self.images.append(f.read())

   def __len__(self):
    return len(self.targets)

   def __getitem__(self, index):
    target = self.targets[index]
    image, retval = cv2.imdecode(self.images[index], cv2.IMREAD_COLOR)
    return image, target

# Profile CPU bottlenecks
python -m cProfile training_script.py --profiling

# Profile GPU bottlenecks
nvprof --print-gpu-trace python train_mnist.py

# Profile system calls bottlenecks
strace -fcT python training_script.py -e trace=open,close,read

Advice 3: *Preprocess everything offline*

• 假设我们为 Cityscapes 训练图像分割模型，其批处理大小为 32，RGB 图像大小是 512x512x3（高、宽、通道）。我们在 CPU 端进行图像标准化（稍后我将会解释为什么这一点比较重要）。在这种情况下，我们最终的图像 tensor 将会是 512 * 512 * 3 * sizeof(float32) = 3,145,728 字节。与批处理大小相乘，结果是 100,663,296 字节，大约 100Mb；

• 除了图像之外，我们还需要提供 ground-truth 掩膜。它们各自的大小为（默认情况下，掩膜的类型是 long，8 个字节）——512 * 512 * 1 * 8 * 32 = 67,108,864 或者大约 67Mb；

• 因此一批数据所需要的总内存是 167Mb。假设有 8 个工作程序，内存的总需求量将是 167 Mb * 8 = 1,336 Mb。

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

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

• GPU 负载不平衡；

• 在主 GPU 上聚合需要额外的视频内存

• 在训练期间继续在前向推导内使用 nn.DataParallel 计算损耗。在这种情况下。za 不会将密集的预测掩码返回给主 GPU，而只会返回单个标量损失；

• 使用分布式训练，也称为 nn.DistributedDataParallel。借助分布式训练的另一个好处是可以看到 GPU 实现 100% 负载。

• medium.com/huggingface/
• medium.com/@theaccelera
• towardsdatascience.com/

def test_loss_profiling():
    loss = nn.BCEWithLogitsLoss()
    with torch.autograd.profiler.profile(use_cuda=True) as prof:
        input = torch.randn((8, 1, 128, 128)).cuda()
        input.requires_grad = True

        target = torch.randint(1, (8, 1, 128, 128)).cuda().float()

        for i in range(10):
            l = loss(input, target)
            l.backward()
    print(prof.key_averages().table(sort_by="self_cpu_time_total"))