Torch的编程方式示例以及必备知识

基于MNIST数据集通过torch实现一个简单的卷积神经网络,也可以用来验证torch是否安装成功

import torch
import torchvision
from torch.autograd import Variable
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
import cv2


## 0. 准备数据
# 导入数据集
train_dataset = datasets.MNIST(root = './data/', train = True, transform = transforms.ToTensor(), download = True)
test_dataset = datasets.MNIST(root = './data/', train = False, transform = transforms.ToTensor(), download = True)
# 数据装载
train_loader = DataLoader(dataset = train_dataset, batch_size=100, shuffle = True)
test_loader = DataLoader(dataset = test_dataset, batch_size=100, shuffle = True)

## 1. 定义模型,当前__init__函数里是集成写法,也可以不采用集成写法,而是离散开,每一层都写一个成员对象self.x
class Model(torch.nn.Module):
    def __init__(self) :
        super(Model, self).__init__()
        self.conv1 = torch.nn.Sequential(torch.nn.Conv2d(1, 64, 3, 1, 1),
                                         torch.nn.ReLU(),
                                         torch.nn.Conv2d(64, 128, 3, 1, 1),
                                         torch.nn.ReLU(),
                                         torch.nn.MaxPool2d(2, 2)) 
        self.dense = torch.nn.Sequential(torch.nn.Linear(14*14*128, 1024),
                                         torch.nn.ReLU(),
                                         torch.nn.Dropout(p = 0.5),
                                         torch.nn.Linear(1024, 10))
        
    def forward(self, x) :
        x = self.conv1(x)
        x = x.view(-1, 14*14*128)
        x = self.dense(x)
        return x

model = Model().to("cuda")
print(model)

# 2. 定义损失函数和优化器
lr = 0.001
cost = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=lr)

# 3. 训练模型
epochs = 5
# tensorboard_path = "./"
# writer = SummaryWriter(tensorboard_path)
model.train()  # 训练模式
num = 0
for epoch in range(epochs) :
    # train
    sum_loss = 0.0      # 统计每个epoch的loss
    train_correct = 0   # 统计每个epoch的correct
    for data in train_loader:
        inputs, labels = data
        # torch0.4后张量与自动微分变量整合,tensor直接当作自动微分变量使用,但Variable仍可使用。
        # inputs, lables = Variable(inputs).cuda(), Variable(labels).cuda()
        inputs, labels = inputs.cuda(), labels.cuda()
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = cost(outputs, labels)
        loss.backward()
        optimizer.step()

        _, id = torch.max(outputs.data, 1)  # Tensor包含求导信息,Tensor.data不包含求导信息
        sum_loss += loss.data               # Tensor.item()转换1*1的张量为数值,Tensor.cpu().numpy()转换张量为ndarray
        train_correct += torch.sum(id == labels.data)

        num += 1
        # writer.add_scalar('Train/Loss', loss.item(), num)
        # writer.add_scalar('Train/Accuracy', torch.sum(id == labels.data).cpu()/id.cpu().numpy().size, num)
        # writer.flush()

    print('[%d,%d] loss:%.03f' % (epoch + 1, epochs, sum_loss / len(train_loader)))
    print('        correct:%.03f%%' % (100 * train_correct / len(train_dataset)))

# 4. 测试模型
model.eval()  # 测试模式
with torch.no_grad():  # 放弃梯度记录节省内存
    test_correct = 0
    for data in test_loader:
        inputs, labels = data
        inputs, labels = inputs.cuda(), labels.cuda()
        outputs = model(inputs)
        _, id = torch.max(outputs.data, 1)
        test_correct += torch.sum(id == labels.data)
    print("correct:%.3f%%" % (100 * test_correct / len(test_dataset)))

torch必备基础知识

torch0.4以后不再使用Variable,在torch中,每一个tensor都包含数据data、梯度grad、求导记录grad_fn,以及requires_grad属性。tensor.data指的就是权重数据本身,但是现在建议更换为tensor.detach(),实际测试发现,torch的官方函数max和min等函数并不要求先获取数据再运算,可以直接送入完整tensor,所以感觉tensor.data用到的场景不多。梯度grad默认为None,只有在运行了loss.backward(),且requires_grad=True时才会有数值。求导记录grad_fn默认是None,在requires_grad=True时,且本身是经过某运算得到的tensor时才会有具体值,用于指示该tensor是如何得到的。还有一点,tensor.item()转换1*1的张量的数据data为fp数值,tensor.item()会自动将数据从GPU移动到CPU。

torch中的tensor能够自动求导,具体规则可以查看【0】。

torch修改现有模型输出

torch官方提供了很多官方模型,那如何修改官方模型的最后一层以满足我们自己的需求呢,当然你可以在定义的时候就指定好输出类别数,下面教程是在你无法这样操作时怎么办。【1】

class VGGNet(nn.Module):
    def __init__(self, num_classes=685):      # num_classes
        super(VGGNet, self).__init__()
        net = models.vgg16(pretrained=True)   # 从预训练模型加载VGG16网络参数
        net.classifier = nn.Sequential()      # 将分类层置空,下面将改变我们的分类层
        self.features = net                   # 保留VGG16的特征层
        self.classifier = nn.Sequential(      # 定义自己的分类层
                nn.Linear(512 * 7 * 7, 1024), # 512 * 7 * 7不能改变 ,由VGG16网络决定的,第二个参数为神经元个数可以微调
                nn.ReLU(True),
                nn.Dropout(0.3),
                nn.Linear(1024, 1024),
                nn.ReLU(True),
                nn.Dropout(0.3),
                nn.Linear(1024, num_classes),
        )

    def forward(self, x):
        x = self.features(x) # 预训练提供的提取特征的部分
        x = x.view(x.size(0), -1)
        x = self.classifier(x) # 自定义的分类部分
        return x
net = VGGNet().to(device)

torch增加现有模型输出

通常在构建网络时,会使用一些比较成熟的网络构建backbone,比如ResNet、MobieNet等等。但有些时候并不需要使用整个backbone,而只需要其中某些层的输出,但自己构建一边backbone又很麻烦。本章主要介绍这种方法就可以很方便地从一个已经搭建好的网络中方便地提取到某些层的输出。【2】与上一章的方法比,本章方法的优势是保留了后半段模型的参数,当然上一章方法也有办法保存后半段参数,所以本章方法可以作为一个扩展。

class IntermediateLayerGetter(nn.ModuleDict):
    """ get the output of certain layers """
    def __init__(self, model, return_layers):
    	# 判断传入的return_layers是否存在于model中
        if not set(return_layers).issubset([name for name, _ in model.named_children()]):
            raise ValueError("return_layers are not present in model")
		
        orig_return_layers = return_layers
        return_layers = {k: v for k, v in return_layers.items()}	# 构造dict
        layers = OrderedDict()
        # 将要从model中获取信息的最后一层之前的模块全部复制下来
        for name, module in model.named_children():
            layers[name] = module
            if name in return_layers:
                del return_layers[name]
            if not return_layers:
                break

        super(IntermediateLayerGetter, self).__init__(layers) # 将所需的网络层通过继承的方式保存下来
        self.return_layers = orig_return_layers

    def forward(self, x):
        out = OrderedDict()
        # 将所需的值以k,v的形式保存到out中
        for name, module in self.named_children():
            x = module(x)
            if name in self.return_layers:
                out_name = self.return_layers[name]
                out[out_name] = x
        return out

import torchvision
    
model = torchvision.models.resnet18()
return_layers = {'layer1':'feature_1', 'layer2':'feature_2'}
backbone = IntermediateLayerGetter(model, return_layers)

backbone.eval()
x = torch.randn(1,3,224,224)
out = backbone(x)
print(out['feature_1'].shape, out['feature_2'].shape)


# torch.Size([1, 64, 56, 56]) torch.Size([1, 128, 28, 28])

torch模型的层访问、模型保存和上载

1. torch模型在定义的时候不能像tf.keras模型一样直接命名,这就给按照名称访问torch模型的某些层带来了麻烦。为了访问某一层,首先需要知道这一层的名称,然后根据名称访问状态字典进而访问到某层的参数。

为了得到每一层的名称,需要遍历所有层并打印名称,从中挑选需要的名称,在torch中遍历所有层有4种方式:

# No.1      [(name, tensor), ...]      <class 'generator'>
parm = {}
for name, parameters in model.named_parameters():
    print(name, "\t", parameters.size())
    parm[name] = parameters.detach().numpy()    # .detach等效于.data, 但是detach更安全建议使用这个

# No.2      {name:tesnor, ...}         <class 'collections.OrderedDict'> 
for param_tensor in model.state_dict():
    print(param_tensor,"\t", model.state_dict()[param_tensor].size())

# No.3      [tensor, ...]              <class 'generator'>
for parameters in model.parameters():
    print(parameters.size())

# No.4      [(name, Sequential), ...]  <class 'generator'>
for name, submodel in model.named_children():
    print(name, type(submodel))
    for name, sublayer in submodel.named_children():
        print(name, sublayer)

# print(type(list(model.named_parameters())[0]))   # ('features.0.0.weight', <class 'torch.nn.parameter.Parameter'>)
# print(model.state_dict()['features.0.0.weight']) # <class 'torch.nn.parameter.Parameter'>
# print(type(list(model.parameters())[0]))         # <class 'torch.nn.parameter.Parameter'>
# print(type(list(model.named_children())[0][1]))  # <class 'torch.nn.modules.container.Sequential'>

第一种方式model.named_parameters()本质上是一个生成器,每次生成的一个数据是一个元组,包含层参数命名和层参数数据,但是这种方式拿到的层参数命名和层参数数据是只是所有可训练参数,并非所有参数,例如BN层的统计均值和方差就访问不到;

第二种方式model.state_dict()本质上是一个有序字典,每一个元素是一个键值对,键是层参数命名和值是层参数数据,这种方式拿到的所有层参数命名和层参数数据,包含BN层的统计均值和方差;

第三种方式model.parameters()本质上是一个生成器,每次生成的一个数据是一个tensor,其就是层参数数据,但是这种方式拿到的层参数数据是只是所有可训练参数,并非所有参数,例如BN层的统计均值和方差就访问不到;

第四种方式model.named_children()本质上是一个生成器,每次生成的一个数据是一个元组,包含子模型或子层命名和子模型或子层(前三种方式访问的都是“层参数”,第四种方式访问的是“层”本体),如果要访问参数,需要在层上使用.parameters()函数【1】;

2. 模型的保存和上载,分保存和上载完整模型(模型结构+模型参数),以及保存和上载模型参数(仅仅指模型参数)。当仅仅保存和上载模型参数时,在保存中其实是先将所有参数变成字典,然后再序列化成.pth模型,在上载中其实是先将.pkl模型反序列化为字典,然后通过字典将参数上载进模型,在这个上载过程要求字典中的键与模型中层参数名称一一对应,字典中的值尺寸与层参数数据尺寸一致。在单纯上载层参数数据时,如果存在命名不一致,需要设定strict=False这个参数【3】,但是依然要求在键能匹配时,值也能匹配。

# No.1 保存模型权重
torch.save(model.state_dict(), "parameter.pth")     # save权重
model.load_state_dict(torch.load('parameter.pth'))  # load

# No.2 保存完整模型
torch.save(model, '\model.pth')                     # 保存完整模型
model = torch.load('\model.pth')

# No.3 根据情况上载模型,不一定上载完整模型
# 读取权重文件,读取为有序字典的形式
assert os.path.exists(weight_path), "file {} dose not exist.".format(weight_path)
pre_weights = torch.load(weight_path, map_location=device)

# 如果最后一层的类别数量相同,则恢复整个模型的权重,如果不相等,则只是恢复前n-1层的权重
pre_dict = {k: v for k, v in pre_weights.items() if model.state_dict()[k].numel() == v.numel()}
missing_keys, unexpected_keys = model.load_state_dict(pre_dict, strict=False)

torch冻结层的方式

torch关于层冻结的方法有两种【4】,但是其实本质上就一种,本质就是设置层参数的属性requires_grad=False,设置 torch.no_grad()最终也是反映在requires_grad属性上,只要层参数的属性requires_grad被设置为False,那在训练的时候,在优化器里即使关联了所有参数也无所谓,但是为了严谨,最好是仅仅关联可训练的参数【5】。在模型推理时,需要使用torch.no_grad(),原因是可以节省内存加快速度。推荐的参数冻结代码如下:

# 冻结n-2层,方式就是设置层的requires_grad为False, model.named_parameters()是列表不包含BN的滑动均值和方差,仅仅包含可训练的参数
freeze_list=list(model.state_dict().keys())[0:-2]   # model.state_dict()是有序字典,包含所有参数
for name, param in model.named_parameters():        # print(model.state_dict()['features.0.1.running_mean'][0:5])
    if name in freeze_list: param.requires_grad=False

# 选择需要训练的参数,其实只要设置了requires_grad属性为false,即使选择所有参数可训练,也不会更新参数了,但是下面写法更科学
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.Adam(params, lr=initial_learning_rate)
# optimizer.param_groups[0]['lr'] = initial_learning_rate   # 更改学习率

torch模型BN层单独失活

在torch中,设置BN层的层参数属性requires_grad=False,只会使得BN层中的两个可训练参数不会更改,并不会影响到统计均值和方差,要使得统计均值和方差不发生变化,可以设置model.eval(),在该模式下依然可以训练,可以更改可训练参数,所以只要模型永远在eval()模式下,可以保证BN层的统计均值和方差不变,但是如果这样设置,那dropout层也永远处在失活状态,这显然不符合要求,所以最好的是单独设置BN层为eval()模式,其它层为train()模式,所以推荐代码如下【6】,在使用model.apply(freeze_bn)时,要搭配model.train()使用,在每次训练时先设定整个模型为train()模式,然后设定BN层为eval()模式:

# 定义函数单独设置所有BN层为推理模式
def freeze_bn(m): 
    if isinstance(m, nn.BatchNorm2d): m.eval() 

model.apply(freeze_bn)

torch模型的transform数据增强

在torch中,通过torchvision库可以实现图像增强,这里只需要说明两点,第一点transform只能接受PIL读取的Image,第二点整个图像增强是在CPU上完成的。当然也可以自定义实现方式,具体见【7】【8】【9】【10】【11】【12】

torch模型的并行化训练和混合精度训练

在torch的分布式并行训练技术中,存在两种方式:第一种方式是DataParallel,这种方式简单易用,只需要修改几行代码即可,但是这种方式只能用在单机多卡,不能用在多机多卡,且存在GPU利用率不均衡问题,甚至有时会存在并行训练速度反而不如单卡训练速度的现象;第二种方式是Distributed DataParallel,这是官方更推荐的使用方式,既可以用在单机多卡,也可以用在多机多卡上,使用起来稍微有点复杂,在教程【13】中有详细介绍,怎么将自己的代码修改为并行训练代码,并在教程【14】中有完整代码。其中,train_sampler.set_epoch(epoch) 必须有这句话,否则每个epoch的数据读取顺序以及在进程中的分布是不变的,具体参考【18】【19】,则读取数据时不需要事先打乱顺序,其会自动打乱,关于num_workers的设置参考【20】,我初始调试代码时将num_workers设置为0,结果无法实现并行训练,修改为8后才可以训练的,还需要说明的是必须在每个进程中都设定完全相同的初始化种子,但是这个相同的种子可以设置为随机初始化,否则每次运行代码,模型的初始化和数据读取顺序是完全相同的。

将自己的代码修改为Distributed DataParallel的并行代码,仅需四步:第一步搭建框架,第二步注册进程,第三步封装模型,第四步封装样本读取器,具体代码如下

# 多机多卡训练模型
import os
from datetime import datetime
import argparse
import torch.multiprocessing as mp
import torchvision
import torchvision.transforms as transforms
import torch
import torch.nn as nn
import torch.distributed as dist
from apex.parallel import DistributedDataParallel as DDP
from apex import amp



# 搭建网络模型
class ConvNet(nn.Module):
    def __init__(self, num_classes=10):
        super(ConvNet, self).__init__()
        self.layer1 = nn.Sequential(
            nn.Conv2d(1, 16, kernel_size=5, stride=1, padding=2),
            nn.BatchNorm2d(16),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2))
        self.layer2 = nn.Sequential(
            nn.Conv2d(16, 32, kernel_size=5, stride=1, padding=2),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2))
        self.fc = nn.Linear(7*7*32, num_classes)

    def forward(self, x):
        out = self.layer1(x)
        out = self.layer2(out)
        out = out.reshape(out.size(0), -1)
        out = self.fc(out)
        return out


## 第一步 搭建框架,相比原始代码增加了一个外包的框架,就是这个main()函数和train()函数
## >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('-n', '--nodes', default=1, type=int, metavar='N', help='number of nodes for distributed training')
    parser.add_argument('-g', '--gpus', default=1, type=int, help='number of gpus per node')
    parser.add_argument('-nr', '--nr', default=0, type=int, help='ranking within the nodes')
    parser.add_argument('--epochs', default=2, type=int, metavar='N', help='number of total epochs to run')
    args = parser.parse_args()                       # -n是几台机子,-g是每台机子上的GPU数量,-nr是机子编号即第几台机子
    args.world_size = args.gpus * args.nodes         # 要建立的总的进程数量,也即总的GPU数量
    os.environ['MASTER_ADDR'] = '10.57.23.164'       # 0号机子的ip地址
    os.environ['MASTER_PORT'] = '8888'               # 0号机子的端口号,随便指定一个未使用的端口号即可
    mp.spawn(train, nprocs=args.gpus, args=(args,))  # 运行起本机对应的所有进程


# 每个进程需要运行的函数,gpu是本机GPU的编号
def train(gpu, args):
    ### 第二步 注册进程,程序知晓了总共有多少进程,当前为第几个进程,以及主机的ip和端口号
    ### >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
    rank = args.nr * args.gpus + gpu                 # 本进程在所有进程的编号
    dist.init_process_group(backend='nccl', init_method='env://', world_size=args.world_size, rank=rank)
    ### <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
    torch.manual_seed(0)                             # 非常重要,必须保证每个进程或GPU上的模型的初始化完全相同
    torch.cuda.manual_seed(0)                        # 但是建议不要设置为定值,可以设置为随机值,先随机化后得到唯一值,在赋值给所有进程
    model = ConvNet()
    torch.cuda.set_device(gpu)
    model.cuda(gpu)
    batch_size = 100
    # define loss function (criterion) and optimizer
    criterion = nn.CrossEntropyLoss().cuda(gpu)
    optimizer = torch.optim.SGD(model.parameters(), 1e-4)
    #### 第三步 封装模型
    #### >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
    model = nn.parallel.DistributedDataParallel(model, device_ids=[gpu])
    #### <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
    # Data loading code
    train_dataset = torchvision.datasets.MNIST(root='./data',
                                               train=True,
                                               transform=transforms.ToTensor(),
                                               download=True)
    ##### 第四步 封装样本读取器
    ##### >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
    train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset,
                                                                    num_replicas=args.world_size,
                                                                    rank=rank)
    ##### <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
    train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
                                               batch_size=batch_size,  # 此batch_size为每个进程所读取数据的size
                                               shuffle=False,          # shuffle必须设置为False,数据顺序会自动打乱,以防万一最好在数据读取阶段就手动打乱顺序
                                               num_workers=0,          # num_workers一般设置大于GPU数量,初始时设置为0,结果程序无法实现并行训练
                                               pin_memory=True,
                                               sampler=train_sampler)  # 当前的方式是数据并行,所以会自动把一个原batch拆分为world_size组,分别送到不同的进程中

    start = datetime.now()
    total_step = len(train_loader)
    for epoch in range(args.epochs):                                         
        train_sampler.set_epoch(epoch)                                 # 必须有这句话,否则每个epoch的数据读取顺序以及在进程中的分布是不变的 
        for i, (images, labels) in enumerate(train_loader):
            images = images.cuda(non_blocking=True)                    # pin_memory=True 与 non_blocking=True搭配使用使得训练变快
            labels = labels.cuda(non_blocking=True)
            # Forward pass
            outputs = model(images)
            loss = criterion(outputs, labels)

            # Backward and optimize
            optimizer.zero_grad()    # 还有一种解决GPU内存不足的办法就是,累计几个batch的梯度后再更新一次
            loss.backward()          # 每个进程基于自己的数据分别求梯度
            optimizer.step()         # 每个进程在获取了其它进程的梯度后求均值,然后更新参数,每个进程中初始化完全相同,更新过程完全相同
            if (i + 1) % 100 == 0 and gpu == 0:
                print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'.format(epoch + 1, args.epochs, i + 1, total_step,
                                                                         loss.item()))
    if gpu == 0:                                                       # 信息输出和模型保存等操作都由编号0的进程操作
        print("Training complete in: " + str(datetime.now() - start))
## <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

## 每个机子上都要运行一遍这个程序,假设总共用4个节点,每个节点有8个gpu
## 0号机子上运行 python src/mnist-distributed.py -n 4 -g 8 -nr 0
## 其余三个机子运行
## python src/mnist-distributed.py -n 4 -g 8 -nr 1
## python src/mnist-distributed.py -n 4 -g 8 -nr 2
## python src/mnist-distributed.py -n 4 -g 8 -nr 3

所谓混合精度训练指的是,模型的参数始终保存为fp32的格式,但是在前向推理时,首先在内存中将fp32格式的权重转换为fp16,图片格式也是fp16,然后进行前向推理,所有中间过程和绝大部分计算全部在fp16格式下进行,最终得到fp16的loss;然后对loss进行放大,然后基于放大后的loss在fp16格式下进行反向梯度计算得到fp16格式下的梯度,然后将fp16格式的梯度转换为fp32格式,然后进行梯度缩小,缩小的倍数与之前loss放大的倍数相同,用于解决fp16的数据溢出问题;然后fp32格式的梯度乘以学习率,然后作用于fp32格式的权重上,至此结束一轮训练,之所以最终权重始终保存为fp32格式是为了解决fp16的舍入误差问题;上面说的前向推理过程中并非所有运算就采用fp16而是部分运算采用fp32是为了提高loss的计算精度。所以模型经过混合精度训练后参数依然是fp32格式,所以此时保存pth,那依然是fp32格式,模型大小不变。之后可以利用其它技术将模型变为fp16格式,理论上经过混合精度训练后生成的fp16模型的精度要高于直接生成的fp16模型。torch进行混合精度训练并不麻烦,教程【17】中提供了完整代码。

torch模型在不同设备之间的转换

一般情况下,torch模型在GPU和CPU之间的转换只需要通过.cuda()和.cpu()就可以进行,或者通过.to(device)实现。但是有些特殊情况,这样无法完成转换,需要手动将每一个张量都转换到目标设备,本条后续需要仔细研究下,当前先做个记录,见【21】

参考

0. PyTorch 的 Autograd

1. Pytorch 构建迁移学习网络——以VGG16,ResNet18和MobileNet V2模型为例

2. Pytorch:提取网络中某些层的输出

3. Pytorch:模型的保存与加载 torch.save()、torch.load()、torch.nn.Module.load_state_dict()

4. pytorch 两种冻结层的方式

5. Pytorch在训练时冻结某些层使其不参与训练(更新梯度)

6. pytorch中的model.eval()和BN层的使用

7. 【pytorch】常用图像处理与数据增强方法合集(torchvision.transforms)

8. 【Pytorch学习笔记】数据增强

9. Pytorch 中的数据增强方式最全解释

10. 图像分类:数据增强(Pytorch版)

11. Pytorch:transforms二十二种数据预处理方法及自定义transforms方法

12. 第九节 学会自定义transforms方法

13. Pytorch中的Distributed Data Parallel与混合精度训练(Apex)

14. 分布式训练代码

15. torch.nn.parallel.DistributedDataParallel: 快速上手

16. pytorch分布式训练(二):torch.nn.parallel.DistributedDataParallel

17. 混合精度训练代码

18. 分布式训练时,torch.utils.data.distributed.DistributedSampler做了什么?

19. Pytorch多机多卡分布式训练​​​​​​​

20. DataLoader的num_workers设置引起卡死/堵塞/报错,以及CPU数据读取加速

21. 通过yolov5训练自己的模型中遇到的一些问题及解决办法​​​​​​​

猜你喜欢

转载自blog.csdn.net/BIT_Legend/article/details/122455702