I have a model trained with 10 epochs and a number of batches less than the total number of batches. My goal is to reload the model and continue training it with the remaining unused batches. I manage to load it but I don't know how to indicate that it will continue training with the rest of the batches. Thanks
def train(model, train_loader, optimizer,scheduler, model_dir, checkpoint_dir, batchEjecutions = 1000):
model.train()
pbar = tqdm(train_loader)
running_loss = 0.0
correct = 0
processed = 0
criterion = nn.CrossEntropyLoss()
for batch_idx, (data, target) in enumerate(pbar):#Recorro cada batch de datos
data, target = data.to(device), target.to(device)
optimizer.zero_grad()
y_pred = model(data)#Predict
loss = criterion(y_pred, target)#Compute de loss functio
running_loss += loss.item()
train_loss.append(loss)
loss.backward()#Backward process ¿In each interction the backward is run?
optimizer.step()
scheduler.step()
pred = y_pred.argmax(dim=1, keepdim=True) # get the index of the max log-probability
correct += pred.eq(target.view_as(pred)).sum().item()
processed += len(data)
#pbar.set_description(desc= f'Loss={loss.item()} Batch_id={batch_idx} Accuracy={100*correct/processed:0.2f} running_loss={running_loss} threshold={best_loss*(0.996)}')
train_acc.append(100*correct/processed)
#pbar.set_description(desc= f'Loss={loss.item()} Batch_id={batch_idx} le={get_lr(optimizer)} Accuracy={100*correct/processed:0.2f}')
if batch_idx >= batchEjecutions:
break
checkpoint = {
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict(),
'schedule': scheduler.state_dict(),
'NBatches': batchEjecutions+1}
save_ckp(checkpoint, True, checkpoint_dir, model_dir)
print(colored('\nTRAINING SUCCESSFULLY COMPLETED', 'cyan'))
import shutil
def save_ckp(state, is_best, checkpoint_dir, model_dir):
torch.save(state, checkpoint_dir)
if is_best:
shutil.copyfile(checkpoint_dir, model_dir)
def load_ckp(checkpoint_fpath, model, optimizer, scheduler):
checkpoint = torch.load(checkpoint_fpath)
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
scheduler.load_state_dict(checkpoint['schedule'])
return model, optimizer, scheduler, checkpoint['epoch']
model=Net7(1.0)
model.to(device)
epochs=5
optimizer = optim.SGD(model.parameters(), lr=0.002, momentum=0.9, weight_decay=9e-4)
scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr=0.02, steps_per_epoch=len(training_generator), pct_start=0.2, div_factor=10, cycle_momentum=False, epochs=epochs)
ckp_path = "/content/drive/MyDrive/CNN_models/checkpoint_model7.pt"
model, optimizer, scheduler, checkpoint_epoch = load_ckp(ckp_path, model, optimizer, scheduler)
Here are my training functions, how I save and reload the model. Now, how do I still train the model with the rest of the barches, if I have only trained it with 150?