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