class Predictor(): """ slim class to make predictions on trained model Args: model - your model (nn.Module) device - the device to send the inputs to callbacks - (optional) list of lpd.callbacks to apply during the predictions phases and states callbacks will be executed by the order of the list name - a friendly identifier Methods: from_trainer - for creating a new Predictor instance from given Trainer from_checkpoint - for creating a new Predictor instance from a saved Trainer checkpoint predict_sample - make prediction on single sample predict_batch - make prediction on single batch predict_data_loader - make prediction on data loader (DataLoader/Iterable/Generator) """ def __init__(self, model, device, callbacks=None, name='lpd predictor'): self._inner = Trainer(model=model, device=device, loss_func=None, optimizer=None, scheduler=None, metrics=None, train_data_loader=None, val_data_loader=None, train_steps=0, val_steps=0, callbacks=callbacks, name=name) @staticmethod def from_trainer(trainer): print(f'[Predictor] - Loading from trainer {trainer.name}') predictor = Predictor(model=trainer.model, device=trainer.device, callbacks=trainer.callbacks, name=f'Predictor-for-Trainer-{trainer.name}') return predictor @staticmethod def from_checkpoint(dir_path, file_name, model, device): full_path = dir_path + file_name checkpoint = T.load(full_path, map_location=device) print(f'[Predictor] - Loading from {full_path}') model.load_state_dict(checkpoint['model']) callbacks = None if 'callbacks' in checkpoint: callbacks = checkpoint['callbacks'] name = None if 'name' in checkpoint: name = checkpoint['name'] predictor = Predictor(model=model, device=device, callbacks=callbacks, name=name) return predictor def predict_sample(self, inputs): return self._inner.predict_sample(inputs) def predict_batch(self, inputs): return self._inner.predict_batch(inputs) def predict_data_loader(self, inputs_data_loader, steps): return self._inner.predict_data_loader(inputs_data_loader, steps)
def test_save_and_predict(self): save_to_dir = os.path.dirname(__file__) + '/trainer_checkpoint/' checkpoint_file_name = 'checkpoint' trainer_file_name = 'trainer' device = tu.get_gpu_device_if_available() model = TestModel().to(device) loss_func = nn.BCEWithLogitsLoss().to(device) optimizer = optim.Adam(model.parameters(), lr=1e-4) scheduler = None metrics = BinaryAccuracyWithLogits(name='acc') callbacks = [ LossOptimizerHandler(), ModelCheckPoint(checkpoint_dir=save_to_dir, checkpoint_file_name=checkpoint_file_name, callback_monitor=CallbackMonitor( monitor_type=MonitorType.LOSS, stats_type=StatsType.VAL, monitor_mode=MonitorMode.MIN), save_best_only=True, save_full_trainer=False), ] data_loader = data_generator() data_loader_steps = 100 num_epochs = 5 trainer = Trainer(model=model, device=device, loss_func=loss_func, optimizer=optimizer, scheduler=scheduler, metrics=metrics, train_data_loader=data_loader, val_data_loader=data_loader, train_steps=data_loader_steps, val_steps=data_loader_steps, callbacks=callbacks, name='Predictor-Trainer-Test') x1_x2, y = next(data_loader) _ = trainer.predict_batch(x1_x2) # JUST TO CHECK THAT IT FUNCTIONS sample = [x1_x2[0][0], x1_x2[1][0]] # PREDICT BEFORE TRAIN sample_prediction_before_train = trainer.predict_sample(sample) trainer.train(num_epochs, verbose=0) # PREDICT AFTER TRAIN sample_prediction_from_trainer = trainer.predict_sample(sample) # SAVE THE TRAINER trainer.save_trainer(save_to_dir, trainer_file_name) #-----------------------------------------------# # CREATE PREDICTOR FROM CURRENT TRAINER #-----------------------------------------------# predictor_from_trainer = Predictor.from_trainer(trainer) # PREDICT FROM PREDICTOR sample_prediction_from_predictor = predictor_from_trainer.predict_sample( sample) self.assertFalse( (sample_prediction_before_train == sample_prediction_from_trainer ).all()) self.assertTrue( (sample_prediction_from_predictor == sample_prediction_from_trainer ).all()) #-----------------------------------------------# # LOAD MODEL CHECKPOINT AS NEW PREDICTOR #-----------------------------------------------# fresh_device = tu.get_gpu_device_if_available() fresh_model = TestModel().to(fresh_device) loaded_predictor = Predictor.from_checkpoint( save_to_dir, checkpoint_file_name + '_best_only', fresh_model, fresh_device) # PREDICT AFTER LOAD sample_prediction_from_loaded_predictor = loaded_predictor.predict_sample( sample) self.assertFalse( (sample_prediction_before_train == sample_prediction_from_trainer ).all()) self.assertTrue((sample_prediction_from_loaded_predictor == sample_prediction_from_trainer).all()) #-----------------------------------------------# # LOAD TRAINER CHECKPOINT AS NEW PREDICTOR #-----------------------------------------------# fresh_device = tu.get_gpu_device_if_available() fresh_model = TestModel().to(fresh_device) loaded_predictor = Predictor.from_checkpoint(save_to_dir, trainer_file_name, fresh_model, fresh_device) # PREDICT AFTER LOAD sample_prediction_from_loaded_predictor = loaded_predictor.predict_sample( sample) self.assertFalse( (sample_prediction_before_train == sample_prediction_from_trainer ).all()) self.assertTrue((sample_prediction_from_loaded_predictor == sample_prediction_from_trainer).all())