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