def train_epoch(epoch, model: BaseModel, dataset, device: str,
tracker: BaseTracker, checkpoint: ModelCheckpoint, log):
model.train()
tracker.reset("train")
train_loader = dataset.train_dataloader()
iter_data_time = time.time()
with Ctq(train_loader) as tq_train_loader:
for i, data in enumerate(tq_train_loader):
data = data.to(device) # This takes time
model.set_input(data)
t_data = time.time() - iter_data_time
iter_start_time = time.time()
model.optimize_parameters(dataset.batch_size)
if i % 10 == 0:
tracker.track(model)
tq_train_loader.set_postfix(**tracker.get_metrics(),
data_loading=float(t_data),
iteration=float(time.time() -
iter_start_time),
color=COLORS.TRAIN_COLOR)
iter_data_time = time.time()
metrics = tracker.publish()
checkpoint.save_best_models_under_current_metrics(model, metrics)
log.info("Learning rate = %f" % model.learning_rate)
def eval_epoch(
epoch: int,
model: BaseModel,
dataset,
device,
tracker: BaseTracker,
checkpoint: ModelCheckpoint,
visualizer: Visualizer,
early_break: bool,
):
model.eval()
tracker.reset("val")
visualizer.reset(epoch, "val")
loader = dataset.val_dataloader()
with Ctq(loader) as tq_val_loader:
for data in tq_val_loader:
data = data.to(device)
with torch.no_grad():
model.set_input(data)
model.forward()
tracker.track(model)
tq_val_loader.set_postfix(**tracker.get_metrics(),
color=COLORS.VAL_COLOR)
if visualizer.is_active:
visualizer.save_visuals(model.get_current_visuals())
if early_break:
break
metrics = tracker.publish(epoch)
tracker.print_summary()
checkpoint.save_best_models_under_current_metrics(model, metrics)
def test_epoch(model: BaseModel, dataset, device, tracker: BaseTracker,
checkpoint: ModelCheckpoint, log):
model.eval()
tracker.reset("test")
loader = dataset.test_dataloader()
with Ctq(loader) as tq_test_loader:
for data in tq_test_loader:
data = data.to(device)
with torch.no_grad():
model.set_input(data)
model.forward()
tracker.track(model)
tq_test_loader.set_postfix(**tracker.get_metrics(),
color=COLORS.TEST_COLOR)
metrics = tracker.publish()
tracker.print_summary()
checkpoint.save_best_models_under_current_metrics(model, metrics)
def train_epoch(
epoch: int,
model: BaseModel,
dataset,
device: str,
tracker: BaseTracker,
checkpoint: ModelCheckpoint,
visualizer: Visualizer,
early_break: bool,
):
model.train()
tracker.reset("train")
visualizer.reset(epoch, "train")
train_loader = dataset.train_dataloader
iter_data_time = time.time()
with Ctq(train_loader) as tq_train_loader:
for i, data in enumerate(tq_train_loader):
model.set_input(data, device)
t_data = time.time() - iter_data_time
iter_start_time = time.time()
model.optimize_parameters(epoch, dataset.batch_size)
if i % 10 == 0:
tracker.track(model)
tq_train_loader.set_postfix(**tracker.get_metrics(),
data_loading=float(t_data),
iteration=float(time.time() -
iter_start_time),
color=COLORS.TRAIN_COLOR)
if visualizer.is_active:
visualizer.save_visuals(model.get_current_visuals())
iter_data_time = time.time()
if early_break:
break
metrics = tracker.publish(epoch)
checkpoint.save_best_models_under_current_metrics(model, metrics)
log.info("Learning rate = %f" % model.learning_rate)
def test_epoch(
epoch: int,
model: BaseModel,
dataset,
device,
tracker: BaseTracker,
checkpoint: ModelCheckpoint,
visualizer: Visualizer,
early_break: bool,
):
model.eval()
loaders = dataset.test_dataloaders()
for idx, loader in enumerate(loaders):
stage_name = dataset.get_test_dataset_name(idx)
tracker.reset(stage_name)
visualizer.reset(epoch, stage_name)
with Ctq(loader) as tq_test_loader:
for data in tq_test_loader:
data = data.to(device)
with torch.no_grad():
model.set_input(data)
model.forward()
tracker.track(model)
tq_test_loader.set_postfix(**tracker.get_metrics(),
color=COLORS.TEST_COLOR)
if visualizer.is_active:
visualizer.save_visuals(model.get_current_visuals())
if early_break:
break
metrics = tracker.publish(epoch)
tracker.print_summary()
checkpoint.save_best_models_under_current_metrics(model, metrics)