当前位置：网站首页>pth 转 onnx 时出现的 gather、unsqueeze 等算子

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

2022-04-23 06:12:00 【wujpbb7】

带动态输入的 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

 

版权声明
本文为[wujpbb7]所创，转载请带上原文链接，感谢
https://blog.csdn.net/blueblood7/article/details/120292177

边栏推荐

猜你喜欢

随机推荐