pth 转 onnx 时出现的 gather、unsqueeze 等算子

带动态输入的 view 或者 reshape 转成 onnx 会有shape/gather/unsqueeze/concat算子。

替换成 flatten 即可。

def forward(self, inputs):
        x1 = self.conv1(inputs)
        x2 = self.conv2(x1)
        # 带动态输入的 view 或者 reshape 转成 onnx 会有shape/gather/unsqueeze/concat算子。
        #x2_flatten = x2.view(x2.size(0), -1)
        #x2_flatten = torch.reshape(x2, (x2.size(0), -1))
        x2_flatten = torch.flatten(x2, start_dim=1)
        x3 = self.fc1(x2_flatten)
        x4 = self.fc2(x3)        
        return x4