Use single model for both training and inference, best way?

Frank_Wessels: <p>Dear readers,</p>
<p>For reasons I would like to train a model (on the GPU) and keep track of “the best” model weights (whenever loss is lowest). Sometimes with the same model I would like to switch to inference and load “the best” weights. Then for reasons when I start training again I would like to start where I left off and <strong>not</strong> at the “best weights”.</p>
<p>So my approach is the following:</p>
<ul>
<li>Saving state dict everytime new lowest loss (deepcopy)</li>
<li>Save (train) state dict everytime when I switch to inference</li>
<li>Load lowest loss state dict when doing inference</li>
<li>Load train state dict when finished with inference</li>
</ul>
<p>It does not seem to work, the best weights state dict keeps getting updated. Deepcopy is not the right approach? I would like to have it all in memory, not saved to a file, for reasons.</p>
<p>Here is it in code:</p>
<p>Trainer:</p>
<pre><code>        loss, model_metrics = self.model(self.batch)
        loss.backward()
        self.optimizer.step()
        self.optimizer.zero_grad()
        self.total_number_of_epochs += 1

        stopping_criterium = model_metrics["loss"]

        if stopping_criterium &lt; self.best_stopping_criterium:
            self.best_stopping_criterium = stopping_criterium
            self.best_loss = model_metrics['loss']

            # Save state dict to disk and to memory
            self.model.save(self.model_weight_path)
            self.best_weights = copy.deepcopy(self.model.model.state_dict())
</code></pre>
<p>Inference:</p>
<pre><code>def inference_single_batch(self, image_batch):
    """
    Inference for a small number of images, should fit completely in (GPU) memory
    """
    # Save current weights (to avoid a infinite loop if predictions are ask every small interval)
    self.train_state_dict = copy.deepcopy(self.model.model.state_dict())

    # Load best model weigths if they exists
    if self.best_weights is not None:
        self.model.model.load_state_dict(self.best_weights)

    # Offload training data if they are on the GPU
    self.training_data_cuda_to_cpu()

    # Set model in evaluate mode if necessary
    if self.model.model.training:
        self.model.model.eval()

    # Perform inference
    image_batch = {Data.IMAGE: torch.from_numpy(np.concatenate(image_batch)).to(self.device)}
    results, _ = self.inference.batch(image_batch)

    # Reload the train weights
    self.model.model.load_state_dict(self.train_state_dict)
</code></pre>
<p>I hope someone can help me,</p>
<p>Regards,<br>
Frank</p>

ptrblck: <aside class="quote no-group" data-username="Frank_Wessels" data-post="1" data-topic="158545">
<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/frank_wessels/48/47093_2.png" class="avatar"> Frank_Wessels:</div>
<blockquote>
<p>It does not seem to work, the best weights state dict keeps getting updated. Deepcopy is not the right approach?</p>
</blockquote>
</aside>
<p>This shouldn’t be the case and <code>deepcopy</code> should work as seen here:</p>
<pre data-code-wrap="python"><code class="lang-python">best_loss = 1000000.
model = nn.Linear(1, 1)
best_state = copy.deepcopy(model.state_dict())

optimizer = torch.optim.Adam(model.parameters(), lr=1.)
criterion = nn.MSELoss()

x = torch.randn(1, 1)
target = torch.randn_like(x)

for epoch in range(10):
    out = model(x)
    loss = criterion(out, target)
    
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()
    
    
    if loss.item() &lt; best_loss:
        print("saving new state_dict")
        print("best_state before\n{}".format(best_state))
        best_loss = loss.item()
        best_state = copy.deepcopy(model.state_dict())
        print('best_state now\n{}'.format(best_state))
</code></pre>

Frank_Wessels: <p>I keep asking stupid late night questions, sorry for that.</p>
<p>Thanks for conformation that it should work. Btw, I did move the saving model before the loss.backward and optimizer stuff, since I thought that is where the weights are updated?</p>
<p>Anyway, the real problem was that I did some augmentation live on the GPU (flipping and stuff) and that when updating the batch I updated the targets but not the data (images) itself. So after “live” annotations (updating by batch data) my result was getting worse and worse and somehow I concluded that since the model for inference with the “best weights” was giving worse results, it had to be from the weights not being right. I underestimated how fast it overfits and how very quickly a worse model can have a lower loss.</p>
<p>Long story short, not sure if you would like to keep my question here, feel free to delete and thanks for the answer (and all other indirect help that I get from you by googling or searching this forum)!</p>