导入
TorchScript 模型追踪与保存
了解 PyTorch 的特性。
将模型保持在指定设备上
1model = Net().to(device)
追踪模型
1trace_input = torch.rand(BATCH_SIZE, IMG_CHANNEL, IMG_HEIGHT, IMG_WIDTH).to(device) 2traced_model = torch.jit.trace(model, trace_input)
将追踪后的模型保存到文件
1traced_model.save("model_jit.pth")