def test_train_step(self):
for model_name in supported_tv_models:
model = cnn.create_vision_cnn(model_name, 10, pretrained=None)
opt = torch.optim.Adam(model.parameters(), lr=1e-3)
loss = nn.CrossEntropyLoss()
# scheduler = torch.optim.lr_scheduler.CyclicLR(optimizer=opt, base_lr=1e-4, max_lr=1e-3, mode="min")
train_metrics = cnn.train_step(model,
train_loader,
loss,
"cpu",
opt,
num_batches=10,
grad_penalty=True)
self.assertIsInstance(train_metrics, Dict)
exp_keys = ("loss", "top1", "top5")
for exp_k in exp_keys:
self.assertTrue(exp_k in train_metrics.keys())
else:
print("Model Created. Training on CPU only")
device = "cpu"
optimizer = optim.Adam(model.parameters(), lr=1e-3)
criterion = (nn.CrossEntropyLoss()
) # All classification problems we need Cross entropy loss
early_stopper = utils.EarlyStopping(patience=7,
verbose=True,
path=SAVE_PATH)
for epoch in tqdm(range(EPOCHS)):
print()
print(f"Training Epoch = {epoch}")
train_metrics = cnn.train_step(model, train_loader, criterion, device,
optimizer)
print()
print(f"Validating Epoch = {epoch}")
valid_metrics = cnn.val_step(model, valid_loader, criterion, device)
validation_loss = valid_metrics["loss"]
early_stopper(validation_loss, model=model)
if early_stopper.early_stop:
print("Saving Model and Early Stopping")
print("Early Stopping. Ran out of Patience for validation loss")
break
print("Done Training, Model Saved to Disk")