I am trying to convert the decoder using torch.jit.script but I am facing some error as below.
My module
import numpy as np
import torch
import torch.nn as nn
from collections import OrderedDict
from layers import *
class DepthDecoder(nn.Module):
def __init__(self, num_ch_enc, scales=range(4), num_output_channels=1, use_skips=True):
super(DepthDecoder, self).__init__()
self.num_output_channels = num_output_channels
self.use_skips = use_skips
self.upsample_mode = 'nearest'
self.scales = scales
self.num_ch_enc = num_ch_enc
self.num_ch_dec = np.array([16, 32, 64, 128, 256])
# decoder
self.convs = OrderedDict()
for i in range(4, -1, -1):
# upconv_0
num_ch_in = self.num_ch_enc[-1] if i == 4 else self.num_ch_dec[i + 1]
num_ch_out = self.num_ch_dec[i]
self.convs[("upconv", i, 0)] = ConvBlock(num_ch_in, num_ch_out)
# upconv_1
num_ch_in = self.num_ch_dec[i]
if self.use_skips and i > 0:
num_ch_in += self.num_ch_enc[i - 1]
num_ch_out = self.num_ch_dec[i]
self.convs[("upconv", i, 1)] = ConvBlock(num_ch_in, num_ch_out)
for s in self.scales:
self.convs[("dispconv", s)] = Conv3x3(self.num_ch_dec[s], self.num_output_channels)
self.decoder = nn.ModuleList(list(self.convs.values()))
self.sigmoid = nn.Sigmoid()
def forward(self, input_features):
outputs = {}
# decoder
x = input_features[-1]
for i in range(4, -1, -1):
x = self.convs[("upconv", i, 0)](x)
x = [upsample(x)]
if self.use_skips and i > 0:
x += [input_features[i - 1]]
x = torch.cat(x, 1)
x = self.convs[("upconv", i, 1)](x)
if i in self.scales:
outputs[("disp", i)] = self.sigmoid(self.convs[("dispconv", i)](x))
return outputs
num_enc_channels = np.array([ 64, 64, 128, 256, 512])
depth_decoder = DepthDecoder( num_ch_enc= num_enc_channels , scales=range(4))
traced_script_module_decoder = torch.jit.script(depth_decoder)
traced_script_module_decoder.save('new-decoder.pt')
Error :
File "C:\Users\lib\site-packages\torch\jit\_recursive.py", line 259, in create_methods_from_stubs
concrete_type._create_methods(defs, rcbs, defaults)
RuntimeError:
Module 'DepthDecoder' has no attribute 'convs' (**This attribute exists on the Python module, but we failed to convert Python type: 'OrderedDict' to a TorchScript type**.):
File "C:\Users\networks\depth_decoder.py", line 55
x = input_features[-1]
for i in range(4, -1, -1):
x = self.convs[("upconv", i, 0)](x)
~~~~~~~~~~ <--- HERE
x = [upsample(x)]
if self.use_skips and i > 0: