pytorch训练加速方案整理

1. 把数据放内存里,降低 io 延迟。

使用:

sudo mount tmpfs /path/to/your/data -t tmpfs -o size=30G

然后把数据放挂载的目录下,即可。

使用时发现应该用一个空目录先挂载内存分区,再把数据拷贝进去。如果直接把数据文件夹挂上去,底下的数据就不可见了。

2. 多GPU并行

首先使用os.environ["CUDA_VISIBLE_DEVICES"] = "0, 1, 2, 3"指定准备使用的GPU

之后使用

Net = net(args)
Net = torch.nn.DataParallel(Net)
Net = Net.to(device)

之后按照正常方法使用即可,需要存权重的话,要用下面这个方法保存

torch.save(model.module.state_dict(), model_out_path)

之前如果存权重的时候没在意这一点,就会报错:

RuntimeError: Error(s) in loading state_dict for ResNet:
Missing key(s) in state_dict: "conv_1.weight", "bn1.weight", "bn1.bias", ......
Unexpected key(s) in state_dict: "module.conv_1.weight", "module.bn1.weight", "module.bn1.bias", ......

解决方法是

# original saved file with DataParallel
state_dict = torch.load(model_path)
# create new OrderedDict that does not contain `module.`
from collections import OrderedDict
new_state_dict = OrderedDict()
for k, v in state_dict.items():
    name = k.replace('.module.','.') # remove `module.`
    new_state_dict[name] = v
# load params
net.load_state_dict(new_state_dict)

ref: https://blog.csdn.net/weixin_41735859/article/details/108610687

3. 使用numba

对于嵌套for循环,尽量将涉及到的变量都转换为numpy的数组(只接受numpy,不能有tensor),将循环放入一个函数里,然后使用numba下面的jit装饰器。例如

from numba import jit

@jit(nopython=True)
def insertion_sort(arr):
    pass