报错
python trans.py
Traceback (most recent call last):
File "trans.py", line 6, in <module>
from nasnet_mobile import nasnetamobile
ModuleNotFoundError: No module named 'nasnet_mobile'
代码
#coding: utf-8
import torch
#import torchvision
# 1.导入PyTorch模型定义
from nasnet_mobile import nasnetamobile
# 2.指定输入大小的shape
dummy_input = torch.randn(1, 3, 224, 224)
# 3. 构建PyTorch model
model = nasnetamobile(121,pretrained=False)
# 4. 载入模型参数
model.load_state_dict(torch.load('L:/Yolov5_DeepSort_Pytorch-master/yolov5/weights/best.pt', map_location='cpu'))
# 5.导出onnx模型文件
torch.onnx.export(model, dummy_input, "nasnet.onnx", verbose=True)
求