Cannot save model

anaandreea1228: <p>Hello!<br>
I am trying to save a model .pt for an inference using torch.jit.trace but I keep getting the same error over and over again. I try to save it so I can use it in libtorch.<br>
I am working on a part on monodepth project where I try to predict disparity on a single image. I am working with the mono_640x192 pretrained model.<br>
Here is the part of the code where I try to save the model:<br>
print("   Loading pretrained encoder")<br>
encoder = networks.ResnetEncoder(18, False)<br>
loaded_dict_enc = torch.load(encoder_path, map_location=device)</p>
<pre><code># extract the height and width of image that this model was trained with
feed_height = loaded_dict_enc['height']
feed_width = loaded_dict_enc['width']
filtered_dict_enc = {k: v for k, v in loaded_dict_enc.items() if k in encoder.state_dict()}
encoder.load_state_dict(filtered_dict_enc)
encoder.to(device)
encoder.eval()

print("   Loading pretrained decoder")
depth_decoder = networks.DepthDecoder(
    num_ch_enc=encoder.num_ch_enc, scales=range(4))

loaded_dict = torch.load(depth_decoder_path, map_location=device)
depth_decoder.load_state_dict(loaded_dict)

depth_decoder.to(device)
depth_decoder.eval()

# FINDING INPUT IMAGES
if os.path.isfile(args.image_path):
    # Only testing on a single image
    paths = [args.image_path]
    output_directory = os.path.dirname(args.image_path)
elif os.path.isdir(args.image_path):
    # Searching folder for images
    paths = glob.glob(os.path.join(args.image_path, '*.{}'.format(args.ext)))
    output_directory = args.image_path
else:
    raise Exception("Can not find args.image_path: {}".format(args.image_path))

print("-&gt; Predicting on {:d} test images".format(len(paths)))

# PREDICTING ON EACH IMAGE IN TURN
with torch.no_grad():
    for idx, image_path in enumerate(paths):

        if image_path.endswith("_disp.jpg"):
            # don't try to predict disparity for a disparity image!
            continue

        # Load image and preprocess
        input_image = pil.open(image_path).convert('RGB')
        original_width, original_height = input_image.size
        input_image = input_image.resize((feed_width, feed_height), pil.LANCZOS)
        input_image = transforms.ToTensor()(input_image).unsqueeze(0)

        # PREDICTION
        input_image = input_image.to(device)
        monodepth_model = MonodepthWrapper(encoder, depth_decoder)

        # Save traced model
        save_traced_model(monodepth_model, input_image, "traced_monodepth_model.pt")
</code></pre>
<p>Here is the wrapper<br>
class MonodepthWrapper(torch.nn.Module):<br>
def <strong>init</strong>(self, encoder, depth_decoder):<br>
super(MonodepthWrapper, self).<strong>init</strong>()<br>
self.encoder = encoder<br>
self.depth_decoder = depth_decoder</p>
<pre><code>def forward(self, x):
    features = self.encoder(x)
    outputs = self.depth_decoder(features)
    return outputs
</code></pre>
<p>Save function<br>
def save_traced_model(model, input_tensor, output_file):<br>
traced_model = torch.jit.trace(model, input_tensor)<br>
torch.jit.save(traced_model, output_file)</p>
<p>Here is the encoder foward function:<br>
def forward(self, input_image):<br>
self.features = <span class="chcklst-box fa fa-square-o fa-fw"></span><br>
x = (input_image - 0.45) / 0.225<br>
x = self.encoder.conv1(x)<br>
x = self.encoder.bn1(x)<br>
self.features.append(self.encoder.relu(x))<br>
self.features.append(self.encoder.layer1(self.encoder.maxpool(self.features[-1])))<br>
self.features.append(self.encoder.layer2(self.features[-1]))<br>
self.features.append(self.encoder.layer3(self.features[-1]))<br>
self.features.append(self.encoder.layer4(self.features[-1]))</p>
<pre><code>    return self.features
</code></pre>
<p>And the decoder one<br>
def forward(self, input_features, *args, **kwargs):<br>
self.outputs = {}</p>
<pre><code>    # 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 &gt; 0:
            x = torch.cat((x, input_features[i - 1]), dim=1)  # Concatenate input_features[i - 1] to tensor x
        x = torch.cat((x,), 1)


        x = self.convs[("upconv", i, 1)](x)
        if i in self.scales:
            self.outputs[("disp", i)] = self.sigmoid(self.convs[("dispconv", i)](x))
    return self.outputs
</code></pre>
<p>This is the error that I get in the console:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/2/f/2fe7ffbe85dd7111d66b774514bc680a4eb09a28.png" data-download-href="https://discuss.pytorch.org/uploads/default/2fe7ffbe85dd7111d66b774514bc680a4eb09a28" title="imageedit_1_2845517073"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/2/f/2fe7ffbe85dd7111d66b774514bc680a4eb09a28_2_690x122.png" alt="imageedit_1_2845517073" data-base62-sha1="6PNuNej13aA3dzoATLMABi5nltS" width="690" height="122" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/2/f/2fe7ffbe85dd7111d66b774514bc680a4eb09a28_2_690x122.png, https://discuss.pytorch.org/uploads/default/original/3X/2/f/2fe7ffbe85dd7111d66b774514bc680a4eb09a28.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/2/f/2fe7ffbe85dd7111d66b774514bc680a4eb09a28.png 2x" data-dominant-color="242424"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">imageedit_1_2845517073</span><span class="informations">1000×178 107 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
What could cause the problem?</p>

ptrblck: <p>I guess this line of code fails:</p>
<pre data-code-wrap="python"><code class="lang-python">self.outputs[("disp", i)] = self.sigmoid(...)
</code></pre>
<p>as the <code>tuple</code> input might not be supported.<br>
Did you try to create a single <code>str</code> as the <code>key</code> for the <code>dict</code>?<br>
Also note that TorchScript is in maintenance mode and the general recommendation is to use <code>torch.compile</code>.</p>

anaandreea1228: <aside class="quote no-group" data-username="ptrblck" data-post="2" data-topic="193492">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>create a single <code>str</code> as the <code>key</code> for the <code>dict</code>?</p>
</blockquote>
</aside>
<p>Thank you! I didn’t try this solution.<br>
I created a single str as the key, but I got another Runtime Error:<br>
RuntimeError: Encountering a dict at the output of the tracer might cause the trace to be incorrect, this is only valid if the container structure does not change based on the module’s inputs. Consider using a constant container instead (e.g. for <code>list</code>, use a <code>tuple</code> instead. for <code>dict</code>, use a <code>NamedTuple</code> instead). If you absolutely need this and know the side effects, pass strict=False to trace() to allow this behavior.<br>
I passed strict = False here</p>
<pre><code class="lang-auto">def save_traced_model(model, input_tensor, output_file):
    traced_model = torch.jit.trace(model, input_tensor, strict=False)
    torch.jit.save(traced_model, output_file)
</code></pre>
<p>and now it works</p>