def main():
parser = argparse.ArgumentParser()
parser.add_argument('--config', required=True)
args = parser.parse_args()
# create logger
logger = get_logger('ModelTrainer')
config = load_config(args.config)
logger.info(config)
# Create model
model = build_model(config)
# Log the number of learnable parameters
logger.info(
f'Number of learnable params {get_number_of_learnable_parameters(model)}'
)
# check if resuming
if config['resume'] is not None:
print(f"Loading checkpoint '{config['resume']}'...")
load_checkpoint(config['resume'], model)
else:
print('start a new training')
# run the training
trainer = BasicFolderTrainer(model, config, logger=logger)
trainer.train()
def __init__(self, model, config, logger=None):
if logger is None:
self.logger = get_logger('ModelTrainer', level=logging.DEBUG)
else:
self.logger = logger
device = config['device']
self.logger.info(f"Sending the model to '{device}'")
self.model = model.to(device)
self.logger.debug(model)
#self.optimizer = optimizer
#self.scheduler = lr_scheduler
#self.loss_criterion = loss_criterion
self.device = device
#self.loaders = loaders
self.config = config