pytorch load state dict_pytorch训练trick

59 阅读 0 评论 39 点赞

我是靠谱客的博主高大石头，最近开发中收集的这篇文章主要介绍pytorch load state dict_pytorch训练trick，觉得挺不错的，现在分享给大家，希望可以做个参考。

概述

pytorch随机种子

pytorch随机种子是随机初始化的，如果想复现一个比较好的结果，可以设置固定随机种子。
其中cudnn打开可以提高计算效率，但是会影响每次复现结果。另外如果图像预处理的时候用了随机切割，随机翻转需要对python的随机数生成器进行固定。

#增加运行效率
torch.backends.cudnn.benchmark = False        # if benchmark=True, deterministic will be False
torch.backends.cudnn.deterministic = True
#固定随机种子
np.random.seed(0)
torch.manual_seed(0)      #cpu设置随机种子
torch.cuda.manual_seed_all(0) #为所有gpu设置随机种子
random.seed(0)

pytorch混合精度训练
pytorch默认用fp32进行训练，如果用fp16做储存和乘法，用fp32做累加计算，fp16会节省显存占用，加快训练和推理的速度。但是也会带来一些问题，具体参考下面链接，所以用混合精度会更好一些。

Nicolas：【PyTorch】唯快不破：基于Apex的混合精度加速zhuanlan.zhihu.com

我一开始用的pip进行安装，发现导入的时候报错，网上解决办法是用源码安装。

pip uninstall apex
git clone https://www.github.com/nvidia/apex
cd apex
python setup.py install

使用方法：

from apex import amp
(net,net_fc),optimizer = amp.initialize([net,net_fc], optimizer, opt_level= "O1")#这里多个net就用列表
with amp.scale_loss(loss_mean, optimizer) as scaled_loss:
scaled_loss.backward()
#bn同步
from apex.parallel import convert_syncbn_model
net = convert_syncbn_model(net)
net = net.to(device)
#多卡
from apex.parallel import DistributedDataParallel as DDP
net, opt = amp.initialize(net, opt, opt_level="O1")  # 字母小写o,不是零。
net = DDP(net, delay_allreduce=True)  # 必须在initialze之后

发现在我电脑上用1050ti对比，3030减小到2725M，速度有一定下降，可能因为我显卡的原因显存变化不明显。
2080ti可以节省一半显存，而且速度单卡混合精度要比之前双卡float32快一点。

多显卡训练

os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"
device_ids = [0,1]
net = nn.DataParallel(net, device_ids=[0,1])#指定显卡
net.to(device=device_ids[0])  #网络放到显卡0

单卡加载多卡训练权重

from collections import OrderedDict
 
state_dictBA = torch.load(opt.pth)
# create new OrderedDict that does not contain `module.`
new_state_dictBA = OrderedDict()
for k, v in state_dictBA.items():
    name = k[7:] # remove `module.`
    new_state_dictBA[name] = v
G_BA.load_state_dict(new_state_dictBA)

tensorboard使用

首先安装 pip tensorboard
在程序里调用，这里只是显示波形图

from torch.utils.tensorboard import SummaryWriter
import time
writer = SummaryWriter('./path/to/log')
for epoch in range(100):
    time.sleep(5)
    print(epoch)
    writer.add_scalar('loss', 100-epoch,epoch)
    writer.add_scalar('acc', epoch/100, epoch)

3、在命令窗口打开界面tensorboard --logdir=log，这里在log路径下打开，否则要指定log文件夹目录。

计算模型耗时：
正确测试代码在cuda运行时间，需要加上torch.cuda.synchronize()，使用该操作来等待GPU全部执行结束，CPU才可以读取时间信息。

torch.cuda.synchronize()#等待当前设备上所有流中的所有核心完成。
start = time.time()
result = model(input)
torch.cuda.synchronize()
end = time.time()

多线程读取数据

当gpu处理数据很快时，cpu读取数据的速度可能会成为瓶颈，导致显卡利用率一直在波动。
可以设置多线程（一般为4,8,16），打开pin_memory(需要内存大一些)

torch.utils.data.DataLoader(image_datasets[x],
                            batch_size=batch_size, 
                            shuffle=True,
                            num_workers=8,
                            pin_memory=True)

半精度推理

在保存模型时，只保存半精度，会节省一般的存储空间，并且在推理时精度不变（我测试时，半精度甚至更好一点点），速度没有测试，理论上应该更快。

参考：

英伟达APEX，多GPU分布式训练，同步Batchnorm，自动混合精度训练法宝指南_我的博客有点东西-CSDN博客_syncbatchnormblog.csdn.net