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