Pytorch gpu加速方法

原文：

https://www.zhihu.com/question/274635237

• relu 用 inplace=True
• 用 eval() 和 with torch.no_grad():
• 每个 batch 后认真的把所有参数从 GPU 拿出来后删除
• 虽然很多回答建议用, 但我建议不要用 torch.cuda.empty_cache() , 这只是释放 GPU 缓存而使得 nvidia-smi 能看得见 pytorch 自动释放的内存而已. 99% 的用户不需要使用这个命令. 并有用户反应每次用反而会减慢 1~2s.[1]
• 注意: 当每张 GPU 里面的 batch_size 太小(<8)时用 batch_norm 会导致训练不稳定, 除非你用以下所说的 APEX 来实现多 GPU sync_bn
• torch.backends.cudnn.deterministic = True 用不用对 GPU 内存占用和效率都没有什么太大的影响. 建议开着.
• 不要用 .cpu() 来取 GPU 里面出来的图片. 这样做的话训练时长可能翻倍.

实现: 研究 pytorch 官方架构就会发现大部分 forward pass 都是 `x = self.conv(x)` 的形式, 很少 introduce new variable. 所以: (1) 把不需要的变量都由 `x` 代替; (2) 变量用完后用 `del` 删除.

例子

  

   

    

     

    
    

     

      def forward(self, x):
     
    
   

    

     

    
    

     

          conv2 = self.conv2(self.conv1(x)) #1/4
     
    
   

    

     

    
    

     
    del x
     
    
   

    

     

    
    

     

          conv3 = self.conv3(conv2) #1/8
     
    
   

    

     

    
    

     

          conv4 = self.conv4(conv3) #1/16
     
    
   

    

     

    
    

     

          conv5 = self.conv5(conv4) #1/32
     
    
   

    

     

    
    

     
 
     
    
   

    

     

    
    

     

          center_64 = self.center_conv1x1(self.center_global_pool(conv5))
     
    
   

    

     

    
    

     
 
     
    
   

    

     

    
    

     

          d5 = self.decoder5(self.center(conv5), conv5)
     
    
   

    

     

    
    

     
    del conv5
     
    
   

    

     

    
    

     

          d4 = self.decoder4(d5, conv4)
     
    
   

    

     

    
    

     
    del conv4
     
    
   

    

     

    
    

     

          d3 = self.decoder3(d4, conv3)
     
    
   

    

     

    
    

     
    del conv3
     
    
   

    

     

    
    

     

          d2 = self.decoder2(d3, conv2)
     
    
   

    

     

    
    

     
    del conv2
     
    
  
 

如果你按照上面的方法把 pin_memory 开启了的话, 请数据放入 GPU 的时候把 non_blocking 开启. 这样如果你只把数据放入 GPU 而不把数据从 GPU 拿出来再做计算的话就会加快很多 (据用户报告可加速 50%). 就算你把 GPU 中数据拿出来 (ie. 用了 .cpu() 命令, 最坏的结果也是与 non_blocking=False 相当:

  

   

    

     

    
    

     

      """Sync Point"""
     
    
   

    

     

    
    

     

      image = image.cuda(non_blocking=True)
     
    
   

    

     

    
    

     

      labels = labels.cuda(non_blocking=True).float()
     
    
   

    

     

    
    

     
 
     
    
   

    

     

    
    

     

      """Async Point"""
     
    
   

    

     

    
    

     

      prediction = net(image)
     
    
  
 

文章来源: blog.csdn.net，作者：AI视觉网奇，版权归原作者所有，如需转载，请联系作者。

原文链接：blog.csdn.net/jacke121/article/details/120190974