当前位置：网站首页>Gather, unsqueeze and other operators when PTH is converted to onnx

Gather, unsqueeze and other operators when PTH is converted to onnx

2022-04-23 07:27:00 【wujpbb7】

With dynamic input view perhaps reshape Turn into onnx There will be shape/gather/unsqueeze/concat operator .

Replace with flatten that will do .

def forward(self, inputs):
        x1 = self.conv1(inputs)
        x2 = self.conv2(x1)
        #  With dynamic input  view  perhaps  reshape  Turn into  onnx  There will be shape/gather/unsqueeze/concat operator .
        #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