代码拉取完成,页面将自动刷新
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()
此处可能存在不合适展示的内容,页面不予展示。您可通过相关编辑功能自查并修改。
如您确认内容无涉及 不当用语 / 纯广告导流 / 暴力 / 低俗色情 / 侵权 / 盗版 / 虚假 / 无价值内容或违法国家有关法律法规的内容,可点击提交进行申诉,我们将尽快为您处理。