Python微信订餐小程序课程视频
https://edu.csdn.net/course/detail/36074
Python实战量化交易理财系统
https://edu.csdn.net/course/detail/35475
定义模型结构
首先使用 PyTorch 定义一个简单的网络模型:
class ConvBnReluBlock(nn.Module):
def \_\_init\_\_(self) -> None:
super().__init__()
self.conv1 = nn.Conv2d(3, 64, 3)
self.bn1 = nn.BatchNorm2d(64)
self.maxpool1 = nn.MaxPool2d(3, 1)
self.conv2 = nn.Conv2d(64, 32, 3)
self.bn2 = nn.BatchNorm2d(32)
self.relu = nn.ReLU()
def forward(self, x):
out = self.conv1(x)
out = self.bn1(out)
out = self.relu(out)
out = self.maxpool1(out)
out = self.conv2(out)
out = self.bn2(out)
out = self.relu(out)
return out
在导出模型之前,需要提前定义一些变量:
model = ConvBnReluBlock() # 定义模型对象
x = torch.randn(2, 3, 255, 255) # 定义输入张量
然后使用 PyTorch 官方 API(torch.onnx.export)导出 ONNX 格式的模型:
# way1:
torch.onnx.export(model, (x), "conv\_bn\_relu\_evalmode.onnx", input_names=["input"], output_names=['output'])
# way