def evaluate(self,
model,
dataloaders,
split,
robust=False,
save_activations=False,
bit_pretrained=False,
adv_metrics=False,
ban_reweight=False):
"""Evaluates the model.
Note:
The latter item in the returned tuple is what is necessary to run
GEORGECluster.train and GEORGECluster.evaluate.
Args:
model(nn.Module): A PyTorch model.
dataloader(DataLoader): The dataloader. The dataset within must
subclass GEORGEDataset.
robust(bool, optional): Whether or not to apply robust optimization. Affects
criterion initialization.
save_activations(bool, optional): If True, saves the activations in
`outputs`. Default is False.
bit_pretrained(bool, optional): If True, assumes bit_pretrained and does not evaluate
performance metrics
Returns:
metrics(Dict[str, Any]) A dictionary object that stores the metrics defined
in self.config['metric_types'].
outputs(Dict[str, Any]) A dictionary object that stores artifacts necessary
for model analysis, including labels, activations, and predictions.
"""
dataloader = dataloaders[split]
# use criterion from training if trained; else, load a new one
if self.criterion is None:
self.criterion = init_criterion(self.config['criterion_config'],
robust, dataloader.dataset,
self.use_cuda)
train_props = np.bincount(
np.array(
dataloaders['train'].dataset.Y_dict['true_subclass'])) / len(
dataloaders['train'].dataset)
split_props = np.bincount(
np.array(dataloader.dataset.Y_dict['true_subclass'])) / len(
dataloader.dataset)
use_cuda = next(model.parameters()).is_cuda
reweight = None if ban_reweight else torch.tensor(train_props /
split_props)
if use_cuda and reweight is not None: reweight = reweight.cuda()
metrics, outputs = self._run_epoch(model,
dataloader,
optimize=False,
save_activations=save_activations,
reweight=reweight,
bit_pretrained=bit_pretrained,
adv_metrics=adv_metrics)
return metrics, outputs
def train(self, model, train_dataloader, val_dataloader, robust=False):
"""Trains the given model.
Note:
Artifacts are only saved if self.save_dir is initialized. Additionally,
this function assumes that the "step" unit of the scheduler is epoch-based.
The model is modified in-place, but the model is also returned to match the
GEORGECluster API.
Args:
model(nn.Module): A PyTorch model.
train_dataloader(DataLoader): The training dataloader. The dataset within must
subclass GEORGEDataset.
val_dataloader(DataLoader): The validation dataloader. The dataset within must
subclass GEORGEDataset.
robust(bool, optional): Whether or not to apply robust optimization. Affects
criterion initialization.
Returns:
model(nn.Module): The best model found during training.
"""
if self.criterion is None:
self.criterion = init_criterion(self.config['criterion_config'],
robust, train_dataloader.dataset,
self.use_cuda)
self.optimizer = init_optimizer(self.config['optimizer_config'], model)
self.scheduler = init_scheduler(self.config['scheduler_config'],
self.optimizer)
# in order to resume model training, load_path must be set explicitly
load_path = self.config.get('load_path', None)
self.state = load_state_dicts(load_path, model, self.optimizer,
self.scheduler, self.logger)
num_epochs = self.config['num_epochs']
checkpoint_metric = self.config['checkpoint_metric']
use_cuda = next(model.parameters()).is_cuda
train_props = np.bincount(
np.array(train_dataloader.dataset.Y_dict['true_subclass'])) / len(
train_dataloader.dataset)
val_props = np.bincount(
np.array(val_dataloader.dataset.Y_dict['true_subclass'])) / len(
val_dataloader.dataset)
reweight = torch.tensor(train_props / val_props)
if use_cuda: reweight = reweight.cuda()
self.logger.basic_info('Starting training.')
for epoch in range(num_epochs):
self.state['epoch'] = epoch
self.scheduler.last_epoch = epoch - 1
self.scheduler.step(
*([self.state[f'best_score']] if type(self.scheduler) ==
schedulers.ReduceLROnPlateau else []))
cur_lr = get_learning_rate(self.optimizer)
self.logger.basic_info(
f'\nEpoch: [{epoch + 1} | {num_epochs}] LR: {cur_lr:.2E}')
self.logger.basic_info('Training:')
train_metrics, _ = self._run_epoch(model,
train_dataloader,
optimize=True,
save_activations=False)
self.logger.basic_info('Validation:')
val_metrics, _ = self._run_epoch(model,
val_dataloader,
optimize=False,
save_activations=False,
reweight=reweight)
metrics = {
**{f'train_{k}': v
for k, v in train_metrics.items()},
**{f'val_{k}': v
for k, v in val_metrics.items()}
}
self._checkpoint(model, metrics, checkpoint_metric, epoch)
self.epoch_logger.append({'learning_rate': cur_lr, **metrics})
if use_cuda: torch.cuda.empty_cache()
best_model_path = os.path.join(self.save_dir, 'best_model.pt')
if os.path.exists(best_model_path):
self.logger.basic_info('\nTraining complete. Loading best model.')
checkpoint = torch.load(best_model_path)
model.load_state_dict(checkpoint['state_dict'])
else:
self.logger.basic_info('Training complete. No best model found.')
return model