pytorch导出torchscript记录

目前网上使用torchscript导出实际模型的案例比较少，以下遇到的问题均是本人导出PatchmatchNet过程中遇到的问题

欢迎使用torchscript导出模型的同学们一起在评论区中讨论遇到的问题，尽量使得本文提出的问题能为我们以后避坑

• 出错提示：aten::div.Tensor(Tensor self, Tensor other) -> (Tensor):
  Expected a value of type ‘Tensor’ for argument ‘self’ but instead found type 'int
  torch.div(feature_

channel, self.G, rounding_mode='trunc')
 整个改成
 int(feature_channel/self.G)

• 出错提示 RuntimeError: undefined value warped_feature
  355行注释掉 del warped_feature, src_feature, src_proj, similarity, view_weight

• 出错提示’int’ object has no attribute or method ‘div_’.:
  similarity = similarity_sum.di

v_(pixel_wise_weight_sum)
改成
similarity = torch.div(similarity_sum, pixel_wise_weight_sum)
或（两个重建的效果不同）
similarity = torch.div(similarity_sum, pixel_wise_weight_sum, rounding_mode='trunc')

patchmatch.py 364行

• Expected a value of type ‘Tensor’ for argument ‘x1’ but instead found type ‘float’
  原因：本来就是tensor的，但是script检测不出来，所以传入前显示的告诉它是tensor转换

score = self.similarity_net(similarity, grid, weight)  # [B, Nde
pth, H, W]
改成
score = self.similarity_net(torch.tensor(similarity), grid, weight)  # [B, Ndepth, H, W]

• torch.jit.frontend.FrontendError: Cannot instantiate class ‘LogSoftmax’ in a script function
  原因：可能真的不支持这个函数吧，但是看了TorchScript — PyTorch 1.13 documentation，里面torchscript对pytorch的支持功能函数是有这个的，最后改用
  import torch.nn.functional as F 用这里面的logsoftmax
  原本使用 import torch.nn 里面的logsoftmax的

softmax = nn.LogSoftmax(dim=1)
score = softmax(score)
改成
score = F.log_softmax(score, dim=1)

• Expected a default value of type Tensor on parameter “grid”.
  原因：可能不能再forward的默认参数列表里面用None

grid: torch.Tensor = None,
weight: torch.Tensor = None,
view_weights: torch.Tensor = None,
改成
grid: torch.Tensor,
weight: torch.Tensor,
view_weights: torch.Tensor,

• RuntimeError:
  Module ‘PatchMatch’ has no attribute ‘propa_conv’ :
  原因：初始判定是init中的控制流导致无法赋值
  原始代码中是一个类的init中有3中不同的初始化，然后放到setattr里面，再在forward里面用getattr根据命名调用，问题就出在这，这个语法torchscript是支持的（查过官方文档），因为我的forward有3个stage，不同的stage对应的conv2d不同，而stage1是没有初始化conv2d的，原始代码中是通过控制流来避免stage1调用conv2d对应的变量，但是导出时一直报错，原因猜测是这个script导出也会把控制流中的内容运行一遍，但是里面的变量没有属性的话就会报上面的错误，因为我的初始化也的确没赋值，最后是把3个stage的初始化变成3个类，代码变长了2倍，没办法，然后不调用conv2d的stage1里面干脆注释掉那段判断代码就好了。
  整段注释掉

还有一个原因就是，torchscript中模型里面的变量，必须是唯一的，不能一个变量对应不同参数的类的实例化，不同参数的类的实例化要用不同名字的变量名来接收

if self.propagate_neighbors > 0:
    # last iteration on stage 1 does not have propagation (photometric consistency filtering)
    # stage 1 的最后一次迭代不需要传播所以无需计算偏移量
    if not (self.stage == 1 and self.patchmatch_iteration == 1):
        propa_offset = self.propa_conv(ref_feature)
        propa_offset = propa_offset.view(batch, 2 * self.propagate_neighbors, height * width)#2是[Δx,Δy]两个方向的偏移
        propa_grid = self.get_propagation_grid(batch, height, width, propa_offset, device, img)

• RuntimeError:
  Class Sigmoid does not have an init function defined:

output = nn.Sigmoid()
x1 = output(x1)
改成
x1 = F.sigmoid(x1)

心得，以后遇到torch.nn的函数当script导出出错时，都用

import torch.nn.functional as F

然后F点出对应的函数，基本都支持的，这样就能通过了，运行测试的结果也是正确的

patchmatch.py 38行

• torch.jit.frontend.FrontendError: Cannot instantiate class ‘Resize’ in a script function:
  这个错误是script里面不支持Resize这个操作，但是代码又必须对图像的尺寸进行修改，所以修改图像尺寸这部分被我封装成一个函数，并将该函数导出成trace模型

self.pre_processing = Pre_processing_img()
改成
self.pre_processing = torch.jit.trace(Pre_processing_img(), torch.rand(1, 5, 3, 480, 640), strict=False)

• need 3 values to unpack but found only 2:
  原因：这个函数里面是根据不同stage返回不同数量的参数，但是C++不能这样动态返回，所以我干脆在stage1那里不对其它stage作判断，不然会编译到别的stage情况下的返回数量

就是返回的数量只能是固定的

depth_sample, score, view_weights = self.evaluation(...)
改成
depth_sample, score = self.evaluation(...)

• Optional [Tuple(tensor, tensor)] canot be used as tuple
  这个问题特别麻烦，一般出现在网络中有的函数需要返回一个tuple的类型
  类似如下

depth_sample, score = self.evaluation(...）

使用script导出时这样是无法返回的，会出现上述报错
我的解决方案，使用trace将该函数单独封装，就是该函数使用trace导出的模型，但是前提是要确保这个函数里面没有if结构
还有就是由于在模型中生成子模型，这个函数默认是在cpu上面的，要添加到gpu上面，否则会出现如下错误提示
Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument mat2 in method wrapper__bmm)
解决方案示例：
注释掉的是原本的

# self.evaluation = Evaluation_1_return2(self.G, self.stage, self.evaluate_neighbors, self.patchmatch_iteration) # 4 1 9 1
src = [torch.rand(1, 16, 240, 320), torch.rand(1, 16, 240, 320), torch.rand(1, 16, 240, 320), torch.rand(1, 16, 240, 320)]
src_p = [torch.rand(1, 4, 4), torch.rand(1, 4, 4), torch.rand(1, 4, 4), torch.rand(1, 4, 4)]
example_input = (
    tocuda(torch.rand(1, 16, 240, 320)),
    tocuda(src),
    tocuda(torch.rand(1, 4, 4)),
    tocuda(src_p),
    tocuda(torch.rand(1, 8, 240, 320)),
    # torch.rand(1),
    # torch.rand(1),
    tocuda(torch.rand(1, 2160, 320, 2)),
    tocuda(torch.rand(1, 8, 9, 240, 320)),
    tocuda(torch.rand(1, 4, 240, 320))
)
self.evaluation = torch.jit.trace(Evaluation_1_return2().cuda(), example_input, strict=False)

提示：如果不用上述方法，就算用以下方法来对返回的Tensor取值也是操作不了的
会报如下错误

• Optional [Tensor] canot be used as Tensor
  还有啥不能索引之类的错误，反正一定要trace那样操作才能解决返回的问题

score = receive[0:8, :, :]
depth_sample = receive[8, :, :]

score, depth_sample = torch.split(receive, [8, 1], dim=0)
depth_sample = receive[0]