上传者: jizhidexiaoming
|
上传时间: 2022-02-12 14:02:41
|
文件大小: 233.08MB
|
文件类型: ONNX
生成该权重代码:
import torch
import torchvision
# use Trace to export onnx model
dummy_input = torch.randn(10, 3, 224, 224, device='cuda') # 定义模型的输入shape
model = torchvision.models.alexnet(pretrained=True).cuda() # if delete cuda(), will generate onnx model with no cuda.
input_names = ['inputs']
output_names = ['outputs']
torch.onnx.export(model, dummy_input, f='alexnet.onnx', verbose=True, input_names=input_names, output_names=output_names, opset_version=10)