Training the saving deep learning model in Pytorch - deep-learning

Question

I have a deep learning model in pytorch (here I provide a simple overview of that). Since I have to run the model each day. I want to save the model in the previous day and then I train the saved model for small number of epochs (3-4) epochs more. Here is the model:
class NET(nn.Module):
def forward(self, y, par):
def test(self, y, num_samples):
def train_batch(model, optimizer, device):
model.train()
optimizer.zero_grad()
length = float(batch.size(0))
...
return loss1
def trainv(model, device, epochs, train_iterator, optimizer, validate_iterator):
for epoch in range(epochs):
for local_batch, local_labels in train_iterator:
train_loss = train_batch(model, optimizer, device)
validate_loss = runs_for_validate(validate_iterator, n_samples) # check on validate data
return train_losses, validate_loss
model = NET(inputs).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr =learning_rate )
train_losses, validate_loss = train(model, device, epochs, train_iterator, optimizer, validate_iterator)
My problem is that I dont know how to save the model and then again loading the saved model to train with more epochs.
previously I could save the model.state_dict() in '.pt' and can load it, but it seems it does not work here. I already saw this post, How can I save my training progress in PyTorch for a certain batch no.? , however I dont know which format I should save the model? and also when I should save the model? Could you please help me with this?

Answer 1

You can save the entire model, including the model architecture and its current state, by passing in the model object to the function.
torch.save(model, "model.pt")
You can save the model's state dictionary, which contains the model's parameters, separately.
torch.save(model.state_dict(), "model_state.pt")
You can load the model by using the torch.load function, passing in the path of the saved model.
model = torch.load("model.pt")
You can also load the state dictionary separately and use it to load the model's parameters.
state_dict = torch.load("model_state.pt")
model.load_state_dict(state_dict)
you can try like
class NET(nn.Module):
def __init__(self, inputs):
super(NET, self).__init__()
self.inputs = inputs
def forward(self, y, par):
# Your forward pass here
pass
def test(self, y, num_samples):
# Your test function here
pass
def train_batch(model, optimizer, device):
model.train()
optimizer.zero_grad()
length = float(batch.size(0))
...
return loss1
def trainv(model, device, epochs, train_iterator, optimizer, validate_iterator, save_path):
best_loss = float('inf')
for epoch in range(epochs):
for local_batch, local_labels in train_iterator:
train_loss = train_batch(model, optimizer, device)
validate_loss = runs_for_validate(validate_iterator, n_samples) # check on validate data
# Save the model if the validation loss improves
if validate_loss < best_loss:
best_loss = validate_loss
torch.save(model.state_dict(), save_path)
return train_losses, validate_loss
def load_model(model, load_path, device):
state_dict = torch.load(load_path, map_location=device)
model.load_state_dict(state_dict)
return model
# Initialize the model
model = NET(inputs).to(device)
# Load the model if a save path is provided
if load_path is not None:
model = load_model(model, load_path, device)
optimizer = torch.optim.Adam(model.parameters(), lr =learning_rate )
train_losses, validate_loss = trainv(model, device, epochs, train_iterator, optimizer, validate_iterator, save_path)