Pytorch distributed 多卡并行载入模型操作

一、Pytorch distributed 多卡并行载入模型

这次来介绍下如何载入模型。

目前没有找到官方的distribute 载入模型的方式,所以采用如下方式。

大部分情况下,我们在测试时不需要多卡并行计算。

所以,我在测试时只使用单卡。

from collections import OrderedDict
device = torch.device("cuda")
model = DGCNN(args).to(device)  #自己的模型
state_dict = torch.load(args.model_path)    #存放模型的位置

new_state_dict = OrderedDict()
for k, v in state_dict.items():
    name = k[7:] # remove `module.`
    new_state_dict[name] = v
    # load params
model.load_state_dict (new_state_dict)

二、pytorch DistributedParallel进行单机多卡训练

One_导入库:

import torch.distributed as dist
from torch.utils.data.distributed import DistributedSampler

Two_进程初始化:

parser = argparse.ArgumentParser()
parser.add_argument('--local_rank', type=int, default=-1)
# 添加必要参数
# local_rank:系统自动赋予的进程编号,可以利用该编号控制打印输出以及设置device

torch.distributed.init_process_group(backend="nccl", init_method='file://shared/sharedfile',
rank=local_rank, world_size=world_size)

# world_size:所创建的进程数,也就是所使用的GPU数量
# (初始化设置详见参考文档)

Three_数据分发:

dataset = datasets.ImageFolder(dataPath)
data_sampler = DistributedSampler(dataset, rank=local_rank, num_replicas=world_size)
# 使用DistributedSampler来为各个进程分发数据,其中num_replicas与world_size保持一致,用于将数据集等分成不重叠的数个子集

dataloader = DataLoader(dataset, batch_size=batch_size, num_workers=1,drop_last=True, pin_memory=True, sampler=data_sampler)
# 在Dataloader中指定sampler时,其中的shuffle必须为False,而DistributedSampler中的shuffle项默认为True,因此训练过程默认执行shuffle

Four_网络模型:

torch.cuda.set_device(local_rank)
device = torch.device('cuda:'+f'{local_rank}')
# 设置每个进程对应的GPU设备

D = Model()
D = torch.nn.SyncBatchNorm.convert_sync_batchnorm(D).to(device)
# 由于在训练过程中各卡的前向后向传播均独立进行,因此无法进行统一的批归一化,如果想要将各卡的输出统一进行批归一化,需要将模型中的BN转换成SyncBN

D = torch.nn.parallel.DistributedDataParallel(
D, find_unused_parameters=True, device_ids=[local_rank], output_device=local_rank)
# 如果有forward的返回值如果不在计算loss的计算图里,那么需要find_unused_parameters=True,即返回值不进入backward去算grad,也不需要在不同进程之间进行通信。

Five_迭代:

data_sampler.set_epoch(epoch)
# 每个epoch需要为sampler设置当前epoch

Six_加载:

dist.barrier()
D.load_state_dict(torch.load('D.pth'), map_location=torch.device('cpu'))
dist.barrier()
# 加载模型前后用dist.barrier()来同步不同进程间的快慢

Seven_启动:

CUDA_VISIBLE_DEVICES=1,3 python -m torch.distributed.launch --nproc_per_node=2 train.py --epochs 15000 --batchsize 10 --world_size 2
# 用-m torch.distributed.launch启动,nproc_per_node为所使用的卡数,batchsize设置为每张卡各自的批大小

以上为个人经验,希望能给大家一个参考,也希望大家多多支持我们。

(0)

相关推荐

  • PyTorch 多GPU下模型的保存与加载(踩坑笔记)

    这几天在一机多卡的环境下,用pytorch训练模型,遇到很多问题.现总结一个实用的做实验方式: 多GPU下训练,创建模型代码通常如下: os.environ['CUDA_VISIBLE_DEVICES'] = args.cuda model = MyModel(args) if torch.cuda.is_available() and args.use_gpu: model = torch.nn.DataParallel(model).cuda() 官方建议的模型保存方式,只保存参数: tor

  • pytorch DistributedDataParallel 多卡训练结果变差的解决方案

    DDP 数据shuffle 的设置 使用DDP要给dataloader传入sampler参数(torch.utils.data.distributed.DistributedSampler(dataset, num_replicas=None, rank=None, shuffle=True, seed=0, drop_last=False)) . 默认shuffle=True,但按照pytorch DistributedSampler的实现: def __iter__(self) -> Ite

  • pytorch多GPU并行运算的实现

    Pytorch多GPU运行 设置可用GPU环境变量.例如,使用0号和1号GPU' os.environ["CUDA_VISIBLE_DEVICES"] = '0,1' 设置模型参数放置到多个GPU上.在pytorch1.0之后的版本中,多GPU运行变得十分方便,先将模型的参数设置并行 if torch.cuda.device_count() > 1: print("Let's use", torch.cuda.device_count(), "GPU

  • Pytorch distributed 多卡并行载入模型操作

    一.Pytorch distributed 多卡并行载入模型 这次来介绍下如何载入模型. 目前没有找到官方的distribute 载入模型的方式,所以采用如下方式. 大部分情况下,我们在测试时不需要多卡并行计算. 所以,我在测试时只使用单卡. from collections import OrderedDict device = torch.device("cuda") model = DGCNN(args).to(device) #自己的模型 state_dict = torch.

  • pytorch fine-tune 预训练的模型操作

    之一: torchvision 中包含了很多预训练好的模型,这样就使得 fine-tune 非常容易.本文主要介绍如何 fine-tune torchvision 中预训练好的模型. 安装 pip install torchvision 如何 fine-tune 以 resnet18 为例: from torchvision import models from torch import nn from torch import optim resnet_model = models.resne

  • 详解如何使用Pytorch进行多卡训练

    目录 1.DP 2.DDP 2.1Pytorch分布式基础 2.2Pytorch分布式训练DEMO 当一块GPU不够用时,我们就需要使用多卡进行并行训练.其中多卡并行可分为数据并行和模型并行.具体区别如下图所示: 由于模型并行比较少用,这里只对数据并行进行记录.对于pytorch,有两种方式可以进行数据并行:数据并行(DataParallel, DP)和分布式数据并行(DistributedDataParallel, DDP). 在多卡训练的实现上,DP与DDP的思路是相似的: 1.每张卡都复制

  • 解决Pytorch 加载训练好的模型 遇到的error问题

    这是一个非常愚蠢的错误 debug的时候要好好看error信息 提醒自己切记好好对待error!切记!切记! -----------------------分割线---------------- pytorch 已经非常友好了 保存模型和加载模型都只需要一条简单的命令 #保存整个网络和参数 torch.save(your_net, 'save_name.pkl') #加载保存的模型 net = torch.load('save_name.pkl') 因为我比较懒我就想直接把整个网络都保存下来,然

  • Pytorch 使用 nii数据做输入数据的操作

    使用pix2pix-gan做医学图像合成的时候,如果把nii数据转成png格式会损失很多信息,以为png格式图像的灰度值有256阶,因此直接使用nii的医学图像做输入会更好一点. 但是Pythorch中的Dataloader是不能直接读取nii图像的,因此加一个CreateNiiDataset的类. 先来了解一下pytorch中读取数据的主要途径--Dataset类.在自己构建数据层时都要基于这个类,类似于C++中的虚基类. 自己构建的数据层包含三个部分 class Dataset(object

  • 解决pytorch多GPU训练保存的模型,在单GPU环境下加载出错问题

    背景 在公司用多卡训练模型,得到权值文件后保存,然后回到实验室,没有多卡的环境,用单卡训练,加载模型时出错,因为单卡机器上,没有使用DataParallel来加载模型,所以会出现加载错误. 原因 DataParallel包装的模型在保存时,权值参数前面会带有module字符,然而自己在单卡环境下,没有用DataParallel包装的模型权值参数不带module.本质上保存的权值文件是一个有序字典. 解决方法 1.在单卡环境下,用DataParallel包装模型. 2.自己重写Load函数,灵活.

  • PyTorch零基础入门之构建模型基础

    目录 一.神经网络的构造 二.神经网络中常见的层 2.1 不含模型参数的层 2.2 含模型参数的层 (1)代码栗子1 (2)代码栗子2 2.3 二维卷积层 stride 2.4 池化层 三.LeNet模型栗子 三点提醒: 四.AlexNet模型栗子 Reference 一.神经网络的构造 PyTorch中神经网络构造一般是基于 Module 类的模型来完成的,它让模型构造更加灵活.Module 类是 nn 模块里提供的一个模型构造类,是所有神经网络模块的基类,我们可以继承它来定义我们想要的模型.

  • pytorch常用函数定义及resnet模型修改实例

    目录 模型定义常用函数 利用nn.Parameter()设计新的层 nn.Sequential nn.ModuleList() nn.ModuleDict() nn.Flatten 模型修改案例 修改模型层 添加外部输入 模型定义常用函数 利用nn.Parameter()设计新的层 import torch from torch import nn class MyLinear(nn.Module): def __init__(self, in_features, out_features):

  • Yii模型操作之criteria查找数据库的方法

    本文实例讲述了Yii模型操作之criteria查找数据库的方法.分享给大家供大家参考,具体如下: 数据模型搜索方法: public function search() { // Warning: Please modify the following code to remove attributes that // should not be searched. $criteria=new CDbCriteria; $criteria->compare('id',$this->id); $c

  • pytorch 加载(.pth)格式的模型实例

    有一些非常流行的网络如 resnet.squeezenet.densenet等在pytorch里面都有,包括网络结构和训练好的模型. pytorch自带模型网址:https://pytorch-cn.readthedocs.io/zh/latest/torchvision/torchvision-models/ 按官网加载预训练好的模型: import torchvision.models as models # pretrained=True就可以使用预训练的模型 resnet18 = mod

随机推荐