Checkpoint Loading not Working As Expected

username: aroibu1
cooked: <p>Hi all! Recently, I started implementing checkpoint saves on my code, as it has become difficult running everything in one go. I save everything as prescribed on the PyTorch tutorials, however when I resume training, the network behaves as if the checkpoint hasn’t been loaded properly. You can see this behaviour in the figure below. I’ll also include relevant extracts from my code, in case anybody can spot if I am doing something silly. Thank you!</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/2/7/2789d2fbcb123257951071779b573b987090151b.png" data-download-href="https://discuss.pytorch.org/uploads/default/2789d2fbcb123257951071779b573b987090151b" title="Screenshot 2020-07-29 at 11.44.31"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/2/7/2789d2fbcb123257951071779b573b987090151b_2_690x163.png" alt="Screenshot 2020-07-29 at 11.44.31" data-base62-sha1="5DLUjgggkVwzwq0SXf5pOb8w6Vl" width="690" height="163" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/2/7/2789d2fbcb123257951071779b573b987090151b_2_690x163.png, https://discuss.pytorch.org/uploads/default/optimized/3X/2/7/2789d2fbcb123257951071779b573b987090151b_2_1035x244.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/2/7/2789d2fbcb123257951071779b573b987090151b_2_1380x326.png 2x" data-dominant-color="FDFCFC"><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">Screenshot 2020-07-29 at 11.44.31</span><span class="informations">1617×382 24.7 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></p>
<pre><code class="lang-auto">

class Solver():    

    def __init__(self ... ):

        self.model = model
        self.optimizer = optimizer(model.parameters(), **optimizer_arguments)
        ....
        self.model_name = model_name
        ....
        self.learning_rate_scheduler = lr_scheduler.StepLR(self.optimizer,
                                                           step_size=learning_rate_scheduler_step_size,
                                                           gamma=learning_rate_scheduler_gamma)
        self.start_epoch = 1
        self.start_iteration = 1
        ....

        if use_last_checkpoint:
            self.load_checkpoint()

    def train(self, train_loader, validation_loader):

        model, optimizer, learning_rate_scheduler = self.model, self.optimizer, self.learning_rate_scheduler

        if torch.cuda.is_available():
            torch.cuda.empty_cache()  # clear memory
            model.cuda(self.device)  # Moving the model to GPU

        iteration = self.start_iteration

        for epoch in range(self.start_epoch, self.number_epochs+1):
            ...... 
                if phase == 'train':
                    model.train()
                else:
                    model.eval()
               .....
                with torch.no_grad():
                    ....
                    if phase == 'validation':
                        early_stop, save_checkpoint = self.EarlyStopping(
                            np.mean(losses))
                        self.early_stop = early_stop
                        if save_checkpoint == True:
                            validation_loss = np.mean(losses)
                            checkpoint_name = os.path.join(
                                self.experiment_directory_path, self.checkpoint_directory, 'checkpoint_epoch_' + str(epoch) + '.' + checkpoint_extension)
                            self.save_checkpoint(state={'epoch': epoch + 1,
                                                        'start_iteration': iteration + 1,
                                                        'arch': self.model_name,
                                                        'state_dict': model.state_dict(),
                                                        'optimizer': optimizer.state_dict(),
                                                        'scheduler': learning_rate_scheduler.state_dict()
                                                        },
                                                 filename=checkpoint_name
                                                 )
                ......
                if phase == 'train':
                    learning_rate_scheduler.step()
    ......

    def save_checkpoint(self, state, filename):

        torch.save(state, filename)

    def load_checkpoint(self, epoch=None):

        ....
            checkpoint_file_path = os.path.join(.....)
            self._checkpoint_reader(checkpoint_file_path)

    def _checkpoint_reader(self, checkpoint_file_path):
       ....
        checkpoint = torch.load(checkpoint_file_path)
        self.start_epoch = checkpoint['epoch']
        self.start_iteration = checkpoint['start_iteration']
        self.model.load_state_dict = checkpoint['state_dict']
        self.optimizer.load_state_dict = checkpoint['optimizer']

        for state in self.optimizer.state.values():
            for key, value in state.items():
                if torch.is_tensor(value):
                    state[key] = value.to(self.device)

        self.learning_rate_scheduler.load_state_dict = checkpoint['scheduler']

</code></pre>

username: aroibu1
cooked: <p>A close-up view from the training performance, without the outliers. As you can see, even after the first peak, the values don’t get close to the point where they were at the stop of training, as if something hasn’t loaded properly. However, when printing values and state_dicts while going through the code step-by-step, all seems to be there.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/6/4/64943436c737f034ff8a9efebafa353cf5078335.png" data-download-href="https://discuss.pytorch.org/uploads/default/64943436c737f034ff8a9efebafa353cf5078335" title="Screenshot 2020-07-29 at 12.18.58"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/6/4/64943436c737f034ff8a9efebafa353cf5078335_2_690x225.png" alt="Screenshot 2020-07-29 at 12.18.58" data-base62-sha1="elLgTKfQkBc41fZyO0SwBQd66UJ" width="690" height="225" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/6/4/64943436c737f034ff8a9efebafa353cf5078335_2_690x225.png, https://discuss.pytorch.org/uploads/default/optimized/3X/6/4/64943436c737f034ff8a9efebafa353cf5078335_2_1035x337.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/6/4/64943436c737f034ff8a9efebafa353cf5078335.png 2x" data-dominant-color="FDF6F4"><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">Screenshot 2020-07-29 at 12.18.58</span><span class="informations">1229×402 59.1 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></p>

username: ptrblck
cooked: <p><code>load_state_dict</code> is a method, which should get the <code>state_dict</code> as its input.<br>
Currently you are reassigning the <code>state_dict</code> to the function:</p>
<pre><code class="lang-python">self.model.load_state_dict = checkpoint['state_dict']
self.optimizer.load_state_dict = checkpoint['optimizer']
</code></pre>
<p>Replace these lines of code with:</p>
<pre><code class="lang-python">self.model.load_state_dict(checkpoint['state_dict'])
self.optimizer.load_state_dict(checkpoint['optimizer'])
</code></pre>
<p>and rerun the code.</p>

username: aroibu1
cooked: <p>Thank you <a class="mention" href="/u/ptrblck">@ptrblck</a>! I’m rerunning now and will let you know of the outcome!</p>