Doubt on In-Memory checkpoint

username: andrejankas
cooked: <p>I would like to save an <strong>in-memory</strong> checkpoint, so to resume training. My question is: Am I saving everything necessary to later continue training with the following code? Also, is it necessary to save the loss?</p>
<pre><code class="lang-python">DEVICE = 'cuda'

def train(model, optimizer, loss, train_loader, test_loader, epochs):
    best_model, best_optimizer, best_loss, best_accuracy = None, None, None, 0.0
    for epoch in range(epochs):
        for x, y in train_loader:
            optimizer.zero_grad()
            x = x.to(DEVICE)
            y = y.to(DEVICE)
            y_hat = model(x)
            output = loss(y_hat, y)
            output.backward()
            optimizer.step()
        model.eval()
        test_acc, test_loss = validate(model, test_loader, loss)
        model.train()
        # Save in-memory checkpoint
        if test_acc &gt; best_accuracy:
            best_accuracy = test_acc
            best_model = copy.deepcopy(model)  # Is this the "correct" way?
            best_optimizer = copy.deepcopy(optimizer)
            best_loss = copy.deepcopy(loss)
    return best_model, best_optimizer, best_loss


def validate(model, test_loader, loss_fn):
    with torch.no_grad():
        pass
</code></pre>
<p>And then, when I call the <strong>train loop</strong>, my intention is to continue training like this:</p>
<pre><code class="lang-python">model = nn.Sequential(
  # ...
)

optimizer = optim.Adam(model.parameters(), lr=1e-3)
loss = nn.CrossEntropyLoss()
model2, optimizer2, loss2 = train(model, optimizer, loss,
                                  train_loader, test_loader, epochs=150)
# Continue training
_, _, _ = train(model2, optimizer2, loss2, train_loader,
                test_loader, epochs=10)
</code></pre>

username: andrejankas
cooked: <p>I am using the same logic I posted here, and when call again the train loop: learning gets stuck, same loss is shown. Any idea why?</p>

username: ptrblck
cooked: <p>I think your code won’t work, since the copied <code>optimizer2</code> will lose its reference to the new copy of <code>model</code> and thus won’t update it anymore.<br>
You would thus have to recreate a new <code>optimizer</code> with <code>model2</code>'s parameter references and load its <code>state_dict</code> as shown here:</p>
<pre><code class="lang-python">def train(model, optimizer, loss, data, target):
    best_model, best_optimizer, best_loss = None, None, None
    
    optimizer.zero_grad()
    out = model(data)
    loss = criterion(out, target)
    loss.backward()
    optimizer.step()

    # Save in-memory checkpoint
    best_model = copy.deepcopy(model)  # Is this the "correct" way?
    best_optimizer = copy.deepcopy(optimizer)
    best_loss = loss.clone()
    return best_model, best_optimizer, best_loss



model = nn.Linear(10, 10)
optimizer = optim.Adam(model.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()
data = torch.randn(1, 10)
target = torch.tensor([0])

model2, optimizer2, loss2 = train(model, optimizer, criterion, data, target)
print(loss2)

for _ in range(10):
    optimizer = optim.Adam(model2.parameters(), lr=1e-3)
    optimizer.load_state_dict(optimizer2.state_dict())
    model2, optimizer2, loss2 = train(model2, optimizer, criterion, data, target)
    print(loss2)
</code></pre>