## PyTorch模型导出为ONNX模型 目前pytorch2paddle主要支持pytorch ScriptModule。 用户可通过如下示例代码,将torchvison或者自己开发写的模型转换成ScriptModule model: ``` #coding: utf-8 import torch import torch.nn as nn from torchvision.models.utils import load_state_dict_from_url # 定义模型 class AlexNet(nn.Module): def __init__(self, num_classes=1000): super(AlexNet, self).__init__() self.features = nn.Sequential( nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2), nn.ReLU(inplace=True), nn.MaxPool2d(kernel_size=3, stride=2), nn.Conv2d(64, 192, kernel_size=5, padding=2), nn.ReLU(inplace=True), nn.MaxPool2d(kernel_size=3, stride=2), nn.Conv2d(192, 384, kernel_size=3, padding=1), nn.ReLU(inplace=True), nn.Conv2d(384, 256, kernel_size=3, padding=1), nn.ReLU(inplace=True), nn.Conv2d(256, 256, kernel_size=3, padding=1), nn.ReLU(inplace=True), nn.MaxPool2d(kernel_size=3, stride=2), ) self.avgpool = nn.AdaptiveAvgPool2d((6, 6)) self.classifier = nn.Sequential( nn.Dropout(0.0), nn.Linear(256 * 6 * 6, 4096), nn.ReLU(inplace=True), nn.Dropout(0.0), nn.Linear(4096, 4096), nn.ReLU(inplace=True), nn.Linear(4096, num_classes), ) def forward(self, x): x = self.features(x) for i in range(1): x = self.avgpool(x) x = torch.flatten(x, 1) x = self.classifier(x) return x # 初始化模型 model = AlexNet() # 加载参数 state_dict = load_state_dict_from_url('https://download.pytorch.org/models/alexnet-owt-4df8aa71.pth', progress=True) model.load_state_dict(state_dict) # 设置模式 model.eval() # 生成ScriptModule并保存 script = torch.jit.script(model) torch.jit.save(script, "alexnet.pt") ```