def train(model, dataloaders: dict, criterion, optimizer, metrics, scheduler, reconstructor, rundir: Union[str, bytes, os.PathLike], stopper, device: torch.device, num_epochs: int = 1000, steps_per_epoch: int = 1000, steps_per_validation_epoch: int = 1000, steps_per_test_epoch: int = 100, early_stopping_begins: int = 0, max_flow: float = 2.5, dali: bool = False, fp16: bool = False): # check our inputs assert (isinstance(model, nn.Module)) assert (isinstance(criterion, nn.Module)) assert (isinstance(optimizer, torch.optim.Optimizer)) scaler = None if fp16: scaler = GradScaler() # loop over number of epochs! for epoch in trange(0, num_epochs): # if our learning rate scheduler plateaus when validation metric saturates, we have to pass our "key metric" for # our validation set. Else, just step every epoch if scheduler.name == 'plateau' and epoch > 0: if hasattr(metrics, 'latest_key'): if 'val' in list(metrics.latest_key.keys()): scheduler.step(metrics.latest_key['val']) elif epoch > 0: scheduler.step() # update the learning rate for this epoch min_lr = utils.get_minimum_learning_rate(optimizer) # store the learning rate for this epoch in our metrics file # print('min lr: {}'.format(min_lr)) metrics.update_lr(min_lr) # loop over our training set! model, metrics, _ = loop_one_epoch(dataloaders['train'], model, criterion, optimizer, metrics, reconstructor, steps_per_epoch, train_mode=True, device=device, dali=dali, fp16=fp16, scaler=scaler) # evaluate on validation set with torch.no_grad(): model, metrics, examples = loop_one_epoch( dataloaders['val'], model, criterion, optimizer, metrics, reconstructor, steps_per_validation_epoch, train_mode=False, device=device, max_flow=max_flow, dali=dali, fp16=fp16, scaler=scaler) # some training protocols do not have test sets, so just reuse validation set for testing inference speed key = 'test' if 'test' in dataloaders.keys() else 'val' loader = dataloaders[key] # evaluate how fast inference takes, without loss calculation, which for some models can have a significant # speed impact metrics = speedtest(loader, model, metrics, steps_per_test_epoch, device=device, dali=dali, fp16=fp16) # use our metrics file to output graphs for this epoch viz.visualize_logger(metrics.fname, examples) # save a checkpoint utils.checkpoint(model, rundir, epoch) # # update latest models file # projects.write_latest_model(config['model'], config['flow_generator'], rundir, config) # input the latest validation loss to the early stopper if stopper.name == 'early': should_stop, _ = stopper(metrics.latest_key['val']) elif stopper.name == 'learning_rate': should_stop = stopper(min_lr) else: # every epoch, increment stopper should_stop = stopper() if should_stop: log.info('Stopping criterion reached!') break return model
def train_from_cfg(cfg: DictConfig) -> Type[nn.Module]: """ train DeepEthogram feature extractors from a configuration object. Args: cfg (DictConfig): configuration object generated by Hydra Returns: trained feature extractor """ rundir = os.getcwd() # done by hydra device = torch.device( "cuda:" + str(cfg.compute.gpu_id) if torch.cuda.is_available() else "cpu") if device != 'cpu': torch.cuda.set_device(device) flow_generator = build_flow_generator(cfg) flow_weights = get_weightfile_from_cfg(cfg, 'flow_generator') assert flow_weights is not None, ( 'Must have a valid weightfile for flow generator. Use ' 'deepethogram.flow_generator.train or cfg.reload.latest') log.info('loading flow generator from file {}'.format(flow_weights)) flow_generator = utils.load_weights(flow_generator, flow_weights, device=device) flow_generator = flow_generator.to(device) dataloaders = get_dataloaders_from_cfg( cfg, model_type='feature_extractor', input_images=cfg.feature_extractor.n_flows + 1) spatial_classifier, flow_classifier = build_model_from_cfg( cfg, return_components=True, pos=dataloaders['pos'], neg=dataloaders['neg']) spatial_classifier = spatial_classifier.to(device) flow_classifier = flow_classifier.to(device) num_classes = len(cfg.project.class_names) utils.save_dict_to_yaml(dataloaders['split'], os.path.join(rundir, 'split.yaml')) criterion = get_criterion(cfg.feature_extractor.final_activation, dataloaders, device) steps_per_epoch = dict(cfg.train.steps_per_epoch) metrics = get_metrics( rundir, num_classes=num_classes, num_parameters=utils.get_num_parameters(spatial_classifier)) dali = cfg.compute.dali # training in a curriculum goes as follows: # first, we train the spatial classifier, which takes still images as input # second, we train the flow classifier, which generates optic flow with the flow_generator model and then classifies # it. Thirdly, we will train the whole thing end to end # Without the curriculum we just train end to end from the start if cfg.feature_extractor.curriculum: del dataloaders # train spatial model, then flow model, then both end-to-end dataloaders = get_dataloaders_from_cfg( cfg, model_type='feature_extractor', input_images=cfg.feature_extractor.n_rgb) log.info('Num training batches {}, num val: {}'.format( len(dataloaders['train']), len(dataloaders['val']))) # we'll use this to visualize our data, because it is loaded z-scored. we want it to be in the range [0-1] or # [0-255] for visualization, and for that we need to know mean and std normalizer = get_normalizer(cfg, input_images=cfg.feature_extractor.n_rgb) optimizer = optim.Adam(filter(lambda p: p.requires_grad, spatial_classifier.parameters()), lr=cfg.train.lr, weight_decay=cfg.feature_extractor.weight_decay) spatialdir = os.path.join(rundir, 'spatial') if not os.path.isdir(spatialdir): os.makedirs(spatialdir) stopper = get_stopper(cfg) # we're using validation loss as our key metric scheduler = initialize_scheduler( optimizer, cfg, mode='min', reduction_factor=cfg.train.reduction_factor) log.info('key metric: {}'.format(metrics.key_metric)) spatial_classifier = train( spatial_classifier, dataloaders, criterion, optimizer, metrics, scheduler, spatialdir, stopper, device, steps_per_epoch, final_activation=cfg.feature_extractor.final_activation, sequence=False, normalizer=normalizer, dali=dali) log.info('Training flow stream....') input_images = cfg.feature_extractor.n_flows + 1 del dataloaders dataloaders = get_dataloaders_from_cfg(cfg, model_type='feature_extractor', input_images=input_images) normalizer = get_normalizer(cfg, input_images=input_images) log.info('Num training batches {}, num val: {}'.format( len(dataloaders['train']), len(dataloaders['val']))) flowdir = os.path.join(rundir, 'flow') if not os.path.isdir(flowdir): os.makedirs(flowdir) flow_generator_and_classifier = FlowOnlyClassifier( flow_generator, flow_classifier).to(device) optimizer = optim.Adam(filter(lambda p: p.requires_grad, flow_classifier.parameters()), lr=cfg.train.lr, weight_decay=cfg.feature_extractor.weight_decay) stopper = get_stopper(cfg) # we're using validation loss as our key metric scheduler = initialize_scheduler( optimizer, cfg, mode='min', reduction_factor=cfg.train.reduction_factor) flow_generator_and_classifier = train( flow_generator_and_classifier, dataloaders, criterion, optimizer, metrics, scheduler, flowdir, stopper, device, steps_per_epoch, final_activation=cfg.feature_extractor.final_activation, sequence=False, normalizer=normalizer, dali=dali) flow_classifier = flow_generator_and_classifier.flow_classifier # overwrite checkpoint utils.checkpoint(flow_classifier, flowdir, stopper.epoch_counter) model = HiddenTwoStream(flow_generator, spatial_classifier, flow_classifier, cfg.feature_extractor.arch, fusion_style=cfg.feature_extractor.fusion, num_classes=num_classes).to(device) # setting the mode to end-to-end would allow to backprop gradients into the flow generator itself # the paper does this, but I don't expect that users would have enough data for this to make sense model.set_mode('classifier') log.info('Training end to end...') input_images = cfg.feature_extractor.n_flows + 1 dataloaders = get_dataloaders_from_cfg(cfg, model_type='feature_extractor', input_images=input_images) normalizer = get_normalizer(cfg, input_images=input_images) optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=cfg.train.lr, weight_decay=cfg.feature_extractor.weight_decay) stopper = get_stopper(cfg) # we're using validation loss as our key metric scheduler = initialize_scheduler( optimizer, cfg, mode='min', reduction_factor=cfg.train.reduction_factor) log.info('Total trainable params: {:,}'.format( utils.get_num_parameters(model))) model = train(model, dataloaders, criterion, optimizer, metrics, scheduler, rundir, stopper, device, steps_per_epoch, final_activation=cfg.feature_extractor.final_activation, sequence=False, normalizer=normalizer, dali=dali) utils.save_hidden_two_stream(model, rundir, dict(cfg), stopper.epoch_counter) return model
def train(model: Type[nn.Module], dataloaders: dict, criterion, optimizer, metrics, scheduler, rundir: Union[str, bytes, os.PathLike], stopper, device: torch.device, steps_per_epoch: dict, final_activation: str = 'sigmoid', sequence: bool = False, class_names: list = None, normalizer=None, dali: bool = False): """ Train feature extractor models Args: model (nn.Module): feature extractor (can also be a component, like the spatial stream or flow stream) dataloaders (dict): dictionary with PyTorch dataloader objects (see dataloaders.py) criterion (nn.Module): loss function optimizer (torch.optim): optimizer (SGD, SGDM, ADAM, etc) metrics (Metrics): metrics object for computing metrics and saving to disk (see metrics.py) scheduler (_LRScheduler): learning rate scheduler (see schedulers.py) rundir (str, os.PathLike): run directory for saving weights stopper (Stopper): object that stops training (see stoppers.py) device (str, torch.device): gpu device steps_per_epoch (dict): keys ['train', 'val', 'test']: number of steps in each "epoch" final_activation (str): either sigmoid or softmax sequence (bool): if True, assumes sequence inputs of shape N,K,T class_names (list): unused normalizer (Normalizer): normalizer object, used for un-zscoring images for visualization purposes Returns: model: a trained model """ # check our inputs assert (isinstance(model, nn.Module)) assert (isinstance(criterion, nn.Module)) assert (isinstance(optimizer, torch.optim.Optimizer)) # loop over number of epochs! for epoch in trange(0, stopper.num_epochs): # if our learning rate scheduler plateaus when validation metric saturates, we have to pass our "key metric" for # our validation set. Else, just step every epoch if scheduler.name == 'plateau' and epoch > 0: if hasattr(metrics, 'latest_key'): if 'val' in list(metrics.latest_key.keys()): scheduler.step(metrics.latest_key['val']) elif epoch > 0: scheduler.step() # update the learning rate for this epoch min_lr = utils.get_minimum_learning_rate(optimizer) # store the learning rate for this epoch in our metrics file # print('min lr: {}'.format(min_lr)) metrics.update_lr(min_lr) # loop over our training set! metrics, _ = loop_one_epoch(dataloaders['train'], model, criterion, optimizer, metrics, final_activation, steps_per_epoch['train'], train_mode=True, device=device, dali=dali) # evaluate on validation set with torch.no_grad(): metrics, examples = loop_one_epoch(dataloaders['val'], model, criterion, optimizer, metrics, final_activation, steps_per_epoch['val'], train_mode=False, sequence=sequence, device=device, normalizer=normalizer, dali=dali) # some training protocols do not have test sets, so just reuse validation set for testing inference speed key = 'test' if 'test' in dataloaders.keys() else 'val' loader = dataloaders[key] # evaluate how fast inference takes, without loss calculation, which for some models can have a significant # speed impact metrics = speedtest(loader, model, metrics, steps_per_epoch['test'], device=device, dali=dali) # use our metrics file to output graphs for this epoch viz.visualize_logger(metrics.fname, examples if len(examples) > 0 else None) # save a checkpoint utils.checkpoint(model, rundir, epoch) # if should_update_latest_models: # projects.write_latest_model(config['model'], config['classifier'], rundir, config) # input the latest validation loss to the early stopper if stopper.name == 'early': should_stop, _ = stopper(metrics.latest_key['val']) elif stopper.name == 'learning_rate': should_stop = stopper(min_lr) else: raise ValueError('Please select a stopping type') if should_stop: log.info('Stopping criterion reached!') break return model
def checkpoint(self, pl_module): utils.checkpoint(pl_module.model, os.getcwd(), pl_module.current_epoch)