Pytorch训练加速技巧

Veröffentlicht am in

Pytorch训练模型，如果数据集比较大的话，需要一块GPU。而如何高效利用这块GPU来，来使得训练速度加快是一件非常重要的是。以下是我实践中的经验和搜索到的一些经验：

  • dataloader中设置num_worker>1, pin_memory=True。num_worker是用于读取数据用的线程数，pin_memory： 是否将数据放在锁页内存，而不是放在缓存区域，因为GPU的显存都在锁页内存中。

    需要注意的是，这个技巧和个人使用的电脑配置有关。在我个人的实验中，通过将num_worker设置为4， pin_memory=True，我训练一个百万图像数据集（224*224）的时间从3个小时缩短到1.5个小时（resnet34为基准）。实际中最好划分出一部分的数据进行“调参”。

  • 使用sgd+momentum这样的需要较少显存的优化器。主要有两个原因，sgd+momentum的效果不比其他的优化器差，反而更能够比Adam达到最优点（即不容易陷在局部极值点和鞍点）。第二个原因是，sgd+momentum只需要对每个参数保存梯度信息和动量信息。而对于类似Adam，额外需要一个二阶信息，这个会加大显存的要求。

  • 使用固态硬盘。我实际中把训练数据放在系统盘，因为系统盘是固态硬盘。把验证集数据放在机械硬盘，毕竟验证过程中可以加大batchsize，且不需要记录梯度信息，显存不会爆炸。

  • 避免将不需要的数据从GPU中拷贝出来。实际中，我们常需要把每个batch的统计信息比如loss，acc或其他数据显示出来。实际上可以在这些数据在GPU中进行累积，在每个epoch结束的 时候显示出来。

  • torch.cuda.emptyCache()：释放PyTorch的缓存分配器中的缓存内存块，用于防止一不小使用大grad_require=True的变量。同时使用del来删除不需要的变量来释放内存

  • 预先读取下一次迭代用的数据：最近才看到一个，mark。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33

class data_prefetcher():
    def __init__(self, loader):
        self.loader = iter(loader)
        # self.stream = torch.cuda.Stream()
        self.preload()

    def preload(self):
        try:
            self.next_data = next(self.loader)
        except StopIteration:
            self.next_input = None
            return
        # with torch.cuda.stream(self.stream):
        #     self.next_data = self.next_data.cuda(non_blocking=True)
            
    def next(self):
        # torch.cuda.current_stream().wait_stream(self.stream)
        data = self.next_data
        self.preload()
        return data

###
正常定义train_loader
###

prefetcher = data_prefetcher(train_loader)
data = prefetcher.next()
i = 0
while data is not None:
    print(i, len(data))
    i += 1
    data = prefetcher.next()