数据处理和数据增强方法
class torchvision.transforms.ToTensor
功能:将 PIL Image 或者 ndarray 转换为 tensor,并且归一化至[0-1]
注意事项:归一化至[0-1]是直接除以 255,若自己的 ndarray 数据尺度有变化,则需要自行
修改。
class torchvision.transforms.Normalize(mean, std)
功能:对数据按通道进行标准化,即先减均值,再除以标准差,注意是 h*w*c
val_transformer = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,))
])
等于模型推理代码中的
input = (input-127.5)/127.5
保存的模型参数实际上一个字典类型,通过key-value的形式来存储模型的所有参数,下面的模型只保存了模型的参数。
if __name__ == "__main__":
model_path = "multi_classifier_22_best.pth"
state_dict = torch.load(model_path, map_location="cuda:1")
print(type(state_dict)) # 类型是 dict
for k,v in state_dict.items():
print(k, v.size())
# features.0.0.weight torch.Size([16, 1, 3, 3])
# features.0.1.weight torch.Size([16])
# features.0.1.bias torch.Size([16])
# features.0.1.running_mean torch.Size([16])
只保存参数
import torch
import torch.nn as nn
class Model(nn.Module):
def __init__(self):
super(Model, self).__init__()
self.fc3 = nn.Linear(10, 1)
def forward(self, x):
x = self.fc3(x)
return x
model = Model()
torch.save(model.state_dict(),"model.pt") # 保存模型参数
模型加载
# 加载模型
model = Model()
weights = torch.load("./model.pth")
model.load_state_dict(weights)
# 如果修改了载入权重或载入权重的结构和当前模型的结构不完全相同,需要加strict=False,这样设置就不需要载入权重和当前网络结构完全匹配
model.load_state_dict(weights, strict=False)
下面官网说发现既可以保存参数,又可以保存网络结构,但是torch.load()必须依赖网络结构定义,还是不行
torch.save(model,"model.pt")
torchscript 可以保存参数和网络结构,需要去验证一下
例子
save_model.py
import torch
import torch.nn as nn
class Model(nn.Module):
def __init__(self):
super(Model, self).__init__()
self.fc3 = nn.Linear(10, 1)
def forward(self, x):
x = self.fc3(x)
return x
model = Model()
torch.save(model,"model.pt") # 保存模型参数
load_model.py
import torch
if __name__ == '__main__':
model_path = "model.pt"
mode = torch.load(model_path)
会报错
AttributeError: Can't get attribute 'Model' on
官网解释
Save and Load the Model — PyTorch Tutorials 1.12.1+cu102 documentation
查看每层对应的名称
for param_tensor in model.state_dict():
print(param_tensor,'\t',model.state_dict()[param_tensor].size())
网路结构和加载模型的权重文件匹配不上
import torch
from collections import OrderedDict
if __name__ == '__main__':
landmark_model = "model.pt"
state_dict = torch.load(landmark_model, map_location="cpu")
# 网路结构和加载模型的权重文件匹配不上,即key对应不上,需要将模型修改成和网络结构一样
new_state_dict = OrderedDict() #将修改好的权重和文件保存到new_state_dict中
for k,v in state_dict.items():
if "module" in k:
name = k[7:]
else:
name = k
new_state_dict[name] = v
# 打印新的模型权重
for k,v in new_state_dict.items():
print(k, v.size())
# 网络结构定义
model = Model()
# 打印网络的state_dict
for k,v in model.state_dict().items():
print(k, v.size())
model.load_state_dict(new_state_dict, strict=True)
# model.load_state_dict(new_state_dict, strict=False) # 如果新的网络结构不一样了,可以将strict改为False
1. RuntimeError: Attempting to deserialize object on CUDA device 3 but torch.cuda.device_count() is 1
原本在服务器上cuda:3上面训练的,然后在本机(本机只有一个显卡)进行测试,出现这种问题,需要在
state_dict = torch.load(model_path)
改为
state_dict = torch.load(model_path, map_location='cuda:0')
# "cuda:1" 表示使用第二块GPU,如果想使用第一块GPU,那么设置为"cuda:0"
device = torch.device("cuda:1")
model = Model()
model = model.to(device) #将模型放在cuda:1上
# 设置成主GPU
device = torch.device("cuda:2")
model = Model()
#device_ids=[2,3] 指定2是主GPU, 3是次GPU,模型和数据由主gpu分发
model = torch.nn.DataParallel(model, device_ids=[2,3])
model = model.to(device)
data = data.to(device) # cuda:2 会将数据平分到各个显卡上的
#注意模型保存
torch.save(model.module.state_dict(), "model.pth")
Ring-Reduce梯度合并 使用DP模式训练保存的模型,然后在自己电脑上运行,需要注意:
DP原理
DP的操作原理是将一个batchsize的输入数据均分到多个GPU上分别计算(此处注意,batchsize要大于GPU个数才能划分)在DP模式中,总共只有一个进程(受到GIL很强限制)。master节点相当于参数服务器,其会向其他卡广播其参数;在梯度反向传播后,各卡将梯度集中到master节点,master节点对搜集来的参数进行平均后更新参数,再将参数统一发送到其他卡上。这种参数更新方式,会导致master节点的计算任务、通讯量很重,从而导致网络阻塞,降低训练速度。
假设读入一个 batch 的数据,其大小为 [30, 5, 2],假设采用三张 GPUs,其运行过程大致为:
将模型放到主 GPU 上,一般为 cuda:0;
把模型同步到 3 张 GPUs 上;
将总输入 batch 的数据平分为 3 份,这里每一份大小为 [10, 5, 2];
依次分别作为每个副本模型的输入;
每个副本模型分别独立进行前向计算,假设为 [4, 5, 2];
从 3 个 GPUs 中收集分别计算后的结果,并按照次序拼接,即 [12, 5, 2],计算 loss;
更新梯度,后向计算.
模型同步 - 数据分发 - 分别前向计算 - loss 计算 - 梯度反传.也就是当你调用nn.DataParallel的时候,只是在你的input数据是并行的,但是你的output loss却不是这样的,每次都会在第一块GPU相加计算,这就造成了第一块GPU的负载远远大于剩余其他的显卡。
主卡显存占用比其他卡会多很多
参考:Pytorch分布式训练/多卡训练(一) —— Data Parallel并行(DP)_hxxjxw的博客-CSDN博客_dataparallel
DDP模式会开启N个进程,每个进程控制一张显卡上加载模型,这些模型相同(被复制了N份到N个显卡),缓解GIL锁的限制。 训练阶段,每个进程通过Ring-Reduce的方法与其他进程通讯(交换各自的梯度) ,使得每个进程都能得到所有梯度之和,各个进程使用平均后的梯度更新自己的参数,因为每个进程下模型的初始参数、更新梯度是一样的,所以更新后模型的参数也保持一致。
Ring-Reduce梯度合并:各个进程独立计算梯度,每个进程将梯度依次传给下一个进程,之后再把从上一个进程拿到的梯度传给下一个进程,循环n(进程数量)次之后,所有的进程就可以得到全部的梯度。(闭环)
相关概念
- rank:用于表示进程的编号/序号(在一些结构图中rank指的是软节点,rank可以看成一个计算单位),每一个进程对应了一个rank的进程,整个分布式由许多rank完成。
- node:物理节点,可以是一台机器也可以是一个容器,节点内部可以有多个GPU。
- rank与local_rank: rank是指在整个分布式任务中进程的序号;local_rank是指在一个node上进程的相对序号,local_rank在node之间相互独立。(注意:在代码中,会使用local_rank来指定GPU,并且local_rank和实际的gpu编号存在映射关系,比如,指定gpu 4,5进行训练,local_rank仍然是0,1,但前提是要先设置os.environ['CUDA_VISIBLE_DEVICES'] = "4,5")。
- nnodes、node_rank与nproc_per_node: nnodes是指物理节点数量,node_rank是物理节点的序号;nproc_per_node是指每个物理节点上面进程的数量。
- word size : 全局(一个分布式任务)中,rank的数量
为了方便理解举个例子,比如分布式中有三台机器,每台机器起4个进程,每个进程占用1个GPU,如下图所示:
注意:
1、rank与GPU之间没有必然的对应关系,一个rank可以包含多个GPU;一个GPU也可以为多个rank服务(多进程共享GPU)。
1. 添加参数和初始化DDP设置
parser = argparse.ArgumentParser()
#使用use_multi_gpu参数控制是否使用多GPU
parser.add_argument('--use_multi_gpu', type = int, default = 0,
help = '')
#设置主GPU,多GPU和单GPU都会使用到
parser.add_argument("--device_id", type=str, default="0")
args = parser.parse_args()
# 如果使用多GPU进行训练
if args.use_multi_gpu == 1:
LOCAL_RANK = int(os.getenv("LOCAL_RANK", -1))
RANK = int(os.getenv("RANK", -1))
WORD_SIZE = int(os.getenv("WORD_SIZE", -1))
torch.distributed.init_process_group(backend="nccl", init_method='env://')
torch.cuda.set_device(LOCAL_RANK)
args.device = torch.device("cuda", LOCAL_RANK)
else:
args.device = torch.device("cuda:{}".format(args.device_ids))
2.数据设置
数据集加载并且设置sampler(通过distributed.DistributedSampler可使得每个gpu训练不同的数据,其内部通过seed种子,设置不同的batch数据),sampler作用:分布式训练中切分数据,是为了让分布式中不同的进程拿不一样的数据,原理是根据进程数量进行切分。
from torch.utils.data.distributed import DistributedSampler
if args.use_multi_gpu == 1:
train_loader = DataLoader(
trainset,
batch_size=args.batch_size,
pin_memory=True,
num_workers=4,
sampler=DistributedSampler(trainset))
else:
train_loader = DataLoader(
trainset,
batch_size=args.batch_size,
shuffle=True,
pin_memory=True,
num_workers=4)
3.模型加载(DDP)
model = model.to(args.device)
if args.use_multi_gpu == 1:
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[LOCAL_RANK], \
output_device=LOCAL_RANK)
4.模型保存
if args.use_multi_gpu == 0:
torch.save(model.state_dict(), model_save_path)
elif args.use_multi_gpu == 1 and LOCAL_RANK == 0: #存储模型时通过local_rank限制单进程写入
torch.save(model.state_dict(), model_save_path)
5. 单机多卡训练
--nproc_per_node 表示使用几个GPU
选择0,3 两个gpu进行训练,
CUDA_VISIBLE_DEVICES="0,3" python3 -m torch.distributed.launch --nproc_per_node=2 train.py \
--use_multi_gpu
选择0,1,2,3 四个gpu进行训练
CUDA_VISIBLE_DEVICES="0,1,2,3" python3 -m torch.distributed.launch --nproc_per_node=4 train.py \
--use_multi_gpu
模型多GPU分布式训练---DDP - 知乎
pytorch分布式多机多卡训练,希望从例子解释,以下代码中参数是什么意思? - 知乎
测试三种方式的训练时间。