生成该权重代码:
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)
2022-02-12 14:02:41
233.08MB
权重备份
1