78 Star 600 Fork 1.2K

Ascend/pytorch

加入 Gitee
与超过 1200万 开发者一起发现、参与优秀开源项目,私有仓库也完全免费 :)
免费加入
文件
克隆/下载
resnet_trace.py 870 Bytes
一键复制 编辑 原始数据 按行查看 历史
import os
import torch
import torchvision
from torch_npu.testing.testcase import TestCase, run_tests
class TestJitTrace(TestCase):
def test_jit_trace_load_script(self):
model = torchvision.models.resnet18().eval().to('npu:0')
example_input = torch.rand(1, 3, 244, 244).to('npu:0')
output = model(example_input)
resnet_model = torch.jit.trace(model, example_input)
torch.jit.save(resnet_model, 'resnet_model.pt')
assert os.path.isfile('./resnet_model.pt')
trace_model = torch.jit.load('./resnet_model.pt')
trace_output = trace_model(example_input)
self.assertRtolEqual(trace_output, output)
script_model = torch.jit.script(model)
script_output = script_model(example_input)
self.assertRtolEqual(script_output, output)
if __name__ == '__main__':
run_tests()
Loading...
马建仓 AI 助手
尝试更多
代码解读
代码找茬
代码优化
Python
1
https://gitee.com/ascend/pytorch.git
git@gitee.com:ascend/pytorch.git
ascend
pytorch
pytorch
master

搜索帮助