How to load pytorch model

kabilan: <p>I have saved my model using the code <code>torch.save(the_model.state_dict(), PATH)</code> after training</p>
<p>while loading, I am confused. can someone explain this code<br>
<code>the_model = TheModelClass(*args, **kwargs) the_model.load_state_dict(torch.load(PATH))</code></p>

ptrblck: <p>The first line of code creates an instance of <code>TheModelClass</code> given the provided arguments.<br>
The parameters of this model will be randomly initialized as was either defined by yourself in the model implementation or just using the default <code>reset_parameter</code> method of the used layers.<br>
Since you’ve stored the <code>state_dict</code> after training (which is the recommended way of serializing the model!), you can now load the trained parameters from the <code>state_dict</code> into your model instance via <code>the_model.load_state_dict(torch.load(path_to_state_dict))</code>.</p>

kabilan: <p>it returns TheModelClass is not defined</p>

ptrblck: <p><code>TheModelClass</code> is just a placeholder name for your model class name.<br>
E.g. you can define your model as:</p>
<pre><code class="lang-python">class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.fc1 = nn.Linear(1, 1)
 
    def forward(self, x):
        return self.fc1(x)

model = MyModel()
model.load_state_dict(torch.load(PATH))
</code></pre>
<p>In that case, you would be using <code>MyModel</code> as the class name instead of <code>TheModelClass</code>.</p>

kabilan: <p>what about *args and **kwargs?</p>

ptrblck: <p>These constructs are used to pass a variable amount of arguments to a class instantiation or function in Python. Have a look at <a href="https://www.geeksforgeeks.org/args-kwargs-python/">this explanation</a> for more information.</p>

alex_d: <p>I was wondering if I could have a clarification on this. For my <code>TheModelClass</code> or <code>MyModel</code> or whatever we name the class, I assume the structure has to match the structure of the original model right?</p>
<p>So for example, if I am fine tuning an fcn segmentation model would the class basically mimic <a href="https://github.com/pytorch/vision/blob/master/torchvision/models/segmentation/_utils.py" rel="nofollow noopener">this structure</a>.</p>
<p>So in my case:</p>
<pre><code class="lang-auto">from collections import OrderedDict

import torch
from torch import nn
from torch.nn import functional as F


class MyModel(nn.Module):
    __constants__ = ['aux_classifier']

    def __init__(self, backbone, classifier, aux_classifier=None):
        super(MyModel, self).__init__()
        self.backbone = backbone
        self.classifier = classifier
        self.aux_classifier = aux_classifier

    def forward(self, x):
        input_shape = x.shape[-2:]
        # contract: features is a dict of tensors
        features = self.backbone(x)

        result = OrderedDict()
        x = features["out"]
        x = self.classifier(x)
        x = F.interpolate(x, size=input_shape, mode='bilinear', align_corners=False)
        result["out"] = x

        if self.aux_classifier is not None:
            x = features["aux"]
            x = self.aux_classifier(x)
            x = F.interpolate(x, size=input_shape, mode='bilinear', align_corners=False)
            result["aux"] = x

        return result
</code></pre>

ptrblck: <p>If you are trying to implement a specific model and would like to keep the workflow, then yes, the structure should be as close as possible.<br>
It seems you would like to write a segmentation model using <code>_SimpleSegmentationModel</code> as the base class. If you derive from this class as your parent class, and want to just reuse e.g. the <code>forward</code> function, you don’t have to copy it and the parent function will be called directly instead.</p>
<p>If you are writing a custom model without any base model, you are free to write it as you wish.</p>

alex_d: <aside class="quote no-group" data-username="ptrblck" data-post="8" data-topic="66432">
<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>It seems you would like to write a segmentation model using <code>_SimpleSegmentationModel</code> as the base class. If you derive from this class as your parent class, and want to just reuse e.g. the <code>forward</code> function, you don’t have to copy it and the parent function will be called directly instead.</p>
</blockquote>
</aside>
<p>You mean do something like this?</p>
<pre><code class="lang-auto">from torch import nn

from ._utils import _SimpleSegmentationModel


__all__ = ["FCN"]


class MyModel(_SimpleSegmentationModel):
    """
    Implements a Fully-Convolutional Network for semantic segmentation.
    Arguments:
        backbone (nn.Module): the network used to compute the features for the model.
            The backbone should return an OrderedDict[Tensor], with the key being
            "out" for the last feature map used, and "aux" if an auxiliary classifier
            is used.
        classifier (nn.Module): module that takes the "out" element returned from
            the backbone and returns a dense prediction.
        aux_classifier (nn.Module, optional): auxiliary classifier used during training
    """
    pass
</code></pre>
<p>I took this from this <a href="https://github.com/pytorch/vision/blob/master/torchvision/models/segmentation/fcn.py" rel="noopener nofollow ugc">page</a></p>
<p>And do I need to have something like this too (taken from the same <a href="https://github.com/pytorch/vision/blob/master/torchvision/models/segmentation/fcn.py" rel="noopener nofollow ugc">page</a>)?</p>
<pre><code class="lang-auto">class FCNHead(nn.Sequential):
    def __init__(self, in_channels, channels):
        inter_channels = in_channels // 4
        layers = [
            nn.Conv2d(in_channels, inter_channels, 3, padding=1, bias=False),
            nn.BatchNorm2d(inter_channels),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Conv2d(inter_channels, channels, 1)
        ]

        super(FCNHead, self).__init__(*layers)
</code></pre>

ptrblck: <p>It really depends on your use case.<br>
Could you explain what you are trying to achieve and based on this we could think about the best approach?</p>
<p>E.g. if you derive your custom model form <code>_SimpleSegmentationModel</code> without changing any method, you could just use the parent model directly.</p>

alex_d: <p>I am just trying to fine tune this model <code>models.segmentation.fcn_resnet101(pretrained=True)</code> with my own data/images to try to get better segmentation accuracy for a specific set of images.</p>

ptrblck: <p>In that case i think just swapping the last layer in the <code>classifier</code> might be sufficient:</p>
<pre><code class="lang-python">model = models.segmentation.fcn_resnet101(pretrained=True)

nb_classes = 2
model.classifier[4] = nn.Conv2d(
    in_channels=512,
    out_channels=nb_classes,
    kernel_size=1,
    stride=1)
</code></pre>
<p>Let me know, if that would work for you or if you need to dig into the model.</p>

alex_d: <p>I’m confused. How would I tie that to the MyModel class?</p>

ptrblck: <p>You wouldn’t need to wrap it in another class or why do you want to use a custom <code>nn.Module</code> on top?</p>

alex_d: <p>From this <a href="https://pytorch.org/tutorials/beginner/saving_loading_models.html" rel="nofollow noopener">tutorial</a> I thought that I had to load the model this way after saving the model.</p>
<p>Save:</p>
<p><code>torch.save(model.state_dict(), PATH)</code></p>
<p>Load:</p>
<pre><code class="lang-auto">model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.eval()
</code></pre>

ptrblck: <p><code>TheModelClass</code> is just a placeholder for any model.<br>
E.g. you could have written a custom model named <code>MyModel</code>, which you could swap then with <code>TheModelClass</code>.</p>
<p>If you want to use a pretrained model, you can directly use it without wrapping into another class.<br>
Of course, you can adapt some layers based on your use case (e.g. changing the last classification layer).</p>

alex_d: <p>So if I understand you correctly in my case I wouldn’t need this line: <code>model = TheModelClass(*args, **kwargs)</code> ?</p>
<p>So just retrieving the model like this would be enough?</p>
<pre><code class="lang-auto">model = torch.load(PATH)
model.load_state_dict(torch.load(PATH))
model.eval()
</code></pre>
<aside class="quote no-group" data-username="ptrblck" data-post="12" data-topic="66432">
<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>In that case i think just swapping the last layer in the <code>classifier</code> might be sufficient:</p>
<pre><code class="lang-auto">model = models.segmentation.fcn_resnet101(pretrained=True)

nb_classes = 2
model.classifier[4] = nn.Conv2d(
    in_channels=512,
    out_channels=nb_classes,
    kernel_size=1,
    stride=1)
</code></pre>
</blockquote>
</aside>
<p>Also, I am a little confused with what you mentioned here, because I already update/ or create a new layer in the finetuning process when I initialize the model like in the function below. So wouldn’t doing this again when retrieving the saved finetuned model be redundant?</p>
<pre><code class="lang-auto">def initialize_model(model_name, num_classes, feature_extract, use_pretrained=True):

    # Initialize these variables which will be set in this if statement. Each of these

    #   variables is model specific.

    model_ft = None

    input_size = 0

    if model_name == "resnet":

        """ FCN_resnet101

        """

        model_ft = models.segmentation.fcn_resnet101(pretrained=use_pretrained)

        set_parameter_requires_grad(model_ft, feature_extract)

        in_chnls = model_ft.classifier[4].in_channels

        model_ft.classifier[4] = nn.Conv2d(in_chnls, num_classes, 1, 1)

        input_size = 750, 1000 

    else:

        print("Invalid model name, exiting...")

        exit()

    return model_ft, input_size
</code></pre>
<p>I then use model_ft as in input in my <code>train_model</code> function which returns the fine tuned model</p>
<pre><code class="lang-auto"># Initialize the model for this run

model_ft, input_size = initialize_model(model_name, num_classes, feature_extract, use_pretrained=True)
</code></pre>
<pre><code class="lang-auto"># Train and evaluate

finetuned_model_ft, hist = train_model(model_ft, dataloaders_dict, criterion, optimizer_ft, num_epochs=num_epochs, is_inception=(model_name=="inception"))
</code></pre>
<p>So the next step would be then to save the new finetuned_model_ft like this:</p>
<p><code>torch.save(finetuned_model_ft.state_dict(), path/to/myModel)</code></p>
<p>and then retrieve it to be used for inference like this:</p>
<pre><code class="lang-auto">my_fine_tuned_model = torch.load(path/to/myModel)
my_fine_tuned_model.load_state_dict(torch.load(path/to/myModel))
my_fine_tuned_model.eval()
</code></pre>
<p>Is this correct?</p>

ptrblck: <p>When restoring the model, I would recommend to recreate your model the same way you did initially, instead of calling <code>my_fine_tuned_model = torch.load(path/to/myModel)</code>.<br>
Besides that, your code should work fine.</p>

alex_d: <p>I’m confused, what do you mean by the same way I did it initially? I mentioned a lot of things in this thread so I’m not sure which way you are talking about haha <img src="https://discuss.pytorch.org/images/emoji/apple/smile.png?v=9" title=":smile:" class="emoji" alt=":smile:"></p>

ptrblck: <p>Haha, sorry for the confusion.<br>
By “initially” I meant you should initialize the model in the same way as you have initialized it before storing the <code>state_dict</code>.</p>
<p>So, if you’ve created the model via <code>model_ft, input_size = initialize_model(...)</code>, and stored the <code>state_dict</code> afterwards, you should use the same approach to recreate your model and then just load the <code>state_dict</code> again.</p>