Torchscript does not take over the weights

Hello together,
I have successfully converted my model to “script”. Only when I test my model now, the weights do not match, so my program does not work. Does anyone have any idea?

import argparse

import os
from os import path
from unifuse import UniFuse
import torch


def load_model(model, args):
    if args.v:
        print("Load Model")
    model_dict = model.state_dict()
    pretrained_dict = torch.load(args.model_path)
    pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
    model_dict.update(pretrained_dict)
    return model

def convert(args):
    model = UniFuse(18, args.height, args.width, True, args.max_depth, "cee", True)
    model = load_model(model, args)
    if args.cuda:
        model.cuda()
        if args.v:
            print("Set model GPU")
    else:
        model.cpu()
        if args.v:
            print("Set model CPU")

    my_script_module = torch.jit.script(model)
    save(my_script_module, args)
    if args.v:
        print(my_script_module.code)

def save(script, args):
    name = "wrapped.pt"
    file = os.path.join(args.output_path, name)
    torch.jit.save(script, file)

def main():
    parser = argparse.ArgumentParser(description="Convert the trained model with torchscript into a binary format")
    # path to the data
    parser.add_argument("--model_path", default="/model.pth", type=str, help="Path to the trained model")
    parser.add_argument("--output_path", default=".", type=str, help="Output path where it will be stored")

    # model params
    parser.add_argument("--num_layers", default=18, type=int, help="Number on layers will be needed")
    parser.add_argument("--height", default=512, type=int, help="Define the height of the input and output Image")
    parser.add_argument("--width", default=1024, type=int, help="Define the width of the input and output")
    parser.add_argument("--max_depth", default=10.0, type=float, help="Define the max depth")
    parser.add_argument("--cuda", default=True, action="store_true", help="if set, use cuda")

    # show the logs
    parser.add_argument("-v", default=True, action="store_true", help="if set, shows the logs")

    args = parser.parse_args()

    if not path.exists(args.model_path):
        exit("Model path didnt exists!")
    if not path.exists(args.output_path):
        exit("Output path didnt exists!")

    convert(args)

if __name__ == "__main__":
    main()

How are you comparing the models after storing them?
Based on your code snippet it seems you are filtering out some keys from the state_dict in load_model? If so, did you checked that this doesn’t create any mismatch in the comparison?

Hey @ptrblck,

Thx for your answer, but the comparison was right. I just forgot the command model.load_state_dict(model_dict) and after this all worked fine.
Thx for your interest and help.