pytorch常用知识记录

1. pytorch常用函数

 1.1 transforms

数据处理和数据增强方法

1.1.1转为 tensor:transforms.ToTensor

class torchvision.transforms.ToTensor
功能:将 PIL Image 或者 ndarray 转换为 tensor,并且归一化至[0-1]
注意事项:归一化至[0-1]是直接除以 255,若自己的 ndarray 数据尺度有变化,则需要自行
修改。

 1.1.2标准化:transforms.Normalize

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

2.模型介绍

2.1 模型文件介绍

保存的模型参数实际上一个字典类型，通过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])

2.2 模型保存和模型加载

只保存参数

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 <module '__main__' from '/home/sunny/code/py/pytorch/load_model.py'>

官网解释 

Save and Load the Model — PyTorch Tutorials 1.12.1+cu102 documentation

2.3 模型文件分析

查看每层对应的名称

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   

2.4 模型加载报错

 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')

3. 单GPU和多GPU训练

3.1 单GPU训练

# "cuda:1" 表示使用第二块GPU，如果想使用第一块GPU，那么设置为"cuda:0"
device = torch.device("cuda:1")

model = Model()
model = model.to(device)  #将模型放在cuda:1上

3.2 DP模式训练

    # 设置成主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模式训练保存的模型，然后在自己电脑上运行，需要注意：

参考：Pytorch分布式训练/多卡训练(一) —— Data Parallel并行(DP)_hxxjxw的博客-CSDN博客_dataparallel

3.3 DDP模式训练

3.3.1 原理

        DDP模式会开启N个进程，每个进程控制一张显卡上加载模型，这些模型相同（被复制了N份到N个显卡），缓解GIL锁的限制。 训练阶段，每个进程通过Ring-Reduce的方法与其他进程通讯（交换各自的梯度） ，使得每个进程都能得到所有梯度之和，各个进程使用平均后的梯度更新自己的参数，因为每个进程下模型的初始参数、更新梯度是一样的，所以更新后模型的参数也保持一致。

Ring-Reduce梯度合并：各个进程独立计算梯度，每个进程将梯度依次传给下一个进程，之后再把从上一个进程拿到的梯度传给下一个进程，循环n（进程数量）次之后，所有的进程就可以得到全部的梯度。（闭环）

3.3.2 实现

相关概念

 为了方便理解举个例子，比如分布式中有三台机器，每台机器起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分布式多机多卡训练，希望从例子解释，以下代码中参数是什么意思？ - 知乎

3.4 速度测试

测试三种方式的训练时间。