pytorch使用GPU运行以及优化方式

Pytorch使用GPU单卡训练

• 1 GPU训练常用代码

• 查看是否本机有cuda可用

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

• 将模型迁移至GPU中进行训练

model = Net(num_classes).to(device)    # 如前面已定义class Net(nn.Module)

• 将输入模型的数据也迁移至GPU中进行训练

images = images.to(device)
labels = labels.to(device)

• 2 训练模型加载与保存

• 模型加载

net = Net()
MODEL_PATH = "./BERT_1.pth"
# 加载模型
net.load_state_dict(torch.load(MODEL_PATH))
# 当CPU加载GPU模型参数时，需要设置map_location='cpu'
net.load_state_dict(torch.load(MODEL_PATH, map_location='cpu'))
# 如果Pytorch的版本不一致会导致模型加载失败
# 解决方法:
# 使用原版本的Pytorch重新加载模型参数并以下面的方式保存
torch.save(net.state_dict(), "BERT_1.pth", _use_new_zipfile_serialization=False)

• 模型保存

torch.save(net.state_dict(), "BERT_1.pth", _use_new_zipfile_serialization=False)

pytorch模型部署

• 1 GPU to CPU

  torch.load('modelparameters.pth’, map_location=lambda storage, loc: storage)

• 2 GPU to GPU 或CPU to CPU

  checkpoint = torch.load('modelparameters.pth’)
  model.load_state_dict(checkpoint)

• 3 GPU0 to GPU1

  torch.load('modelparameters.pth’, map_location={'cuda:1’:'cuda:0’})

• 4 CPU to GPU1

  torch.load('modelparameters.pth’, map_location=lambda storage, loc: storage.cuda(1))

  CPU模型部署优化

• 所谓的CPU优化, 工业场景下基本都放在预测阶段.(因为用CPU训练模型, 也就谈不上什么硬件优化了)

• 1 模型量化(当前只支持在cpu上进行,1.3.0以上版本)

• 通俗的理解, 就是将模型的参数精度进行降低操作, 用更少的比特位(torch.qint8)代替较多的比特位(torch.float32), 从而缩减模型, 并加速推断速度.

注意：量化会损耗模型的精度。生成式模型一般不使用量化，分类模型使用量化。

model.load_state_dict(torch.load(origin_model_path))
print(model)
print('-------------------------------------------------------------------')
# 使用torch.quantization.quantize_dynamic获得动态量化的模型
# 量化的网络层为所有的nn.Linear的权重，使其成为int8
quantized_model_1 = torch.quantization.quantize_dynamic(model, {torch.nn.Linear}, dtype=torch.qint8)

• 2 模型大小打印

# 打印模型大小
def print_size_of_model(model):
    # 保存模型中的参数部分到持久化文件
    torch.save(model.state_dict(), "temp.p")
    # 打印持久化文件的大小
    print('Size (MB):', os.path.getsize("temp.p")/1e6)
    # 移除该文件
    os.remove('temp.p')