如何在pytorch中部署半精度模型

如何在pytorch中部署半精度模型？很多新手对此不是很清楚，为了帮助大家解决这个难题，下面小编将为大家详细讲解，有这方面需求的人可以来学习下，希望你能有所收获。

具体方法

在pytorch中，一般模型定义都继承torch.nn.Moudle，torch.nn.Module基类的half()方法会把所有参数转为16位浮点，所以在模型加载后，调用一下该方法即可达到模型切换的目的.接下来只需要在推断时把input的tensor切换为16位浮点即可

另外还有一个小的trick，在推理过程中模型输出的tensor自然会成为16位浮点，如果需要新创建tensor，最好调用已有tensor的new_zeros，new_full等方法而不是torch.zeros和torch.full，前者可以自动继承已有tensor的类型，这样就不需要到处增加代码判断是使用16位还是32位了，只需要针对input tensor切换.

补充：pytorch 使用amp.autocast半精度加速训练

准备工作

pytorch 1.6+

如何使用autocast？

根据官方提供的方法，

答案就是autocast + GradScaler。

1，autocast

正如前文所说，需要使用torch.cuda.amp模块中的autocast 类。使用也是非常简单的：

如何在PyTorch中使用自动混合精度？

答案：autocast + GradScaler。

1.autocast

正如前文所说，需要使用torch.cuda.amp模块中的autocast 类。使用也是非常简单的

from torch.cuda.amp import autocast as autocast

# 创建model，默认是torch.FloatTensor
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)

for input, target in data:
    optimizer.zero_grad()

    # 前向过程(model + loss)开启 autocast
    with autocast():
        output = model(input)
        loss = loss_fn(output, target)

    # 反向传播在autocast上下文之外
    loss.backward()
    optimizer.step()

2.GradScaler

GradScaler就是梯度scaler模块，需要在训练最开始之前实例化一个GradScaler对象。

因此PyTorch中经典的AMP使用方式如下：

from torch.cuda.amp import autocast as autocast

# 创建model，默认是torch.FloatTensor
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)
# 在训练最开始之前实例化一个GradScaler对象
scaler = GradScaler()

for epoch in epochs:
    for input, target in data:
        optimizer.zero_grad()

        # 前向过程(model + loss)开启 autocast
        with autocast():
            output = model(input)
            loss = loss_fn(output, target)

        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()

3.nn.DataParallel

单卡训练的话上面的代码已经够了，亲测在2080ti上能减少至少1/3的显存，至于速度。。。

要是想多卡跑的话仅仅这样还不够，会发现在forward里面的每个结果都还是float32的，怎么办？

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()

    def forward(self, input_data_c1):
     with autocast():
      # code
     return

pytorch的优点

1.PyTorch是相当简洁且高效快速的框架；2.设计追求最少的封装；3.设计符合人类思维，它让用户尽可能地专注于实现自己的想法；4.与google的Tensorflow类似，FAIR的支持足以确保PyTorch获得持续的开发更新；5.PyTorch作者亲自维护的论坛 供用户交流和求教问题6.入门简单

看完上述内容是否对您有帮助呢？如果还想对相关知识有进一步的了解或阅读更多相关文章，请关注亿速云行业资讯频道，感谢您对亿速云的支持。