X, y = torch.rand(num_samples, num_features), torch.rand(num_samples) dataset = TensorDataset(X, y) loader = DataLoader(dataset, batch_size=32, num_workers=1) loaders = {"train": loader, "valid": loader} # model, criterion, optimizer, scheduler model = torch.nn.Linear(num_features, 1) criterion = torch.nn.MSELoss() optimizer = torch.optim.Adam(model.parameters()) scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, [3, 6]) runner = SupervisedRunner() # model training runner.train( model=model, criterion=criterion, optimizer=optimizer, scheduler=scheduler, loaders=loaders, logdir="./logdir", num_epochs=8, verbose=True, check=True, load_best_on_end=True, ) # model inference for prediction in runner.predict_loader(loader=loader): assert prediction["logits"].cpu().detach().numpy().shape == (32, 1) # model tracing traced_model = runner.trace(loader=loader)
from src.model.mobilenet import MBv2 from src.model.model_wrapper import ModelWrapper if __name__ == "__main__": image_size = [1, 3, 416, 416] batch = {'image': np.random.randn(*image_size)} # create model model = MBv2(num_classes=2) model_wrapper = ModelWrapper(model) device = torch.device("cpu") print(f"device: {device}") logdir = "logs/segmentation" checkpoint_path = f"{logdir}/checkpoints/best.pth" checkpoint = utils.load_checkpoint(f"{logdir}/checkpoints/best.pth") model_wrapper.load_state_dict(checkpoint['model_state_dict']) # create runner runner = SupervisedRunner(device=device, input_key="image", input_target_key="mask") # trace model # saves to `logdir` and returns a `ScriptModule` class traced_script_module = runner.trace(model=model_wrapper, batch=batch, fp16=False) traced_script_module.save("traced_model.pt")
if i >= max_count: break # ## Model tracing # # Catalyst allows you to use Runner to make [tracing](https://pytorch.org/docs/stable/jit.html) models. # # > How to do this in the Config API, we wrote in [our blog (issue \#2)](https://github.com/catalyst-team/catalyst-info#catalyst-info-2-tracing-with-torchjit) # # For this purpose it is necessary to pass in a method `trace ` model and a batch on which `predict_batch ` will be executed: # In[ ]: batch = next(iter(loaders["valid"])) # saves to `logdir` and returns a `ScriptModule` class runner.trace(model=model, batch=batch, logdir=logdir, fp16=is_fp16_used) get_ipython().system('ls {logdir}/trace/') # After this, you can easily load the model and predict anything! # In[ ]: if is_fp16_used: model = utils.load_traced_model( f"{logdir}/trace/traced-forward-opt_O1.pth", device="cuda", opt_level="O1") else: model = utils.load_traced_model(f"{logdir}/trace/traced-forward.pth", device="cpu")
def train_segmentation_model( model: torch.nn.Module, logdir: str, num_epochs: int, loaders: Dict[str, DataLoader] ): criterion = { "dice": DiceLoss(), "iou": IoULoss(), "bce": nn.BCEWithLogitsLoss() } learning_rate = 0.001 encoder_learning_rate = 0.0005 layerwise_params = {"encoder*": dict(lr=encoder_learning_rate, weight_decay=0.00003)} model_params = utils.process_model_params(model, layerwise_params=layerwise_params) base_optimizer = RAdam(model_params, lr=learning_rate, weight_decay=0.0003) optimizer = Lookahead(base_optimizer) scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, factor=0.25, patience=2) device = utils.get_device() runner = SupervisedRunner(device=device, input_key='image', input_target_key='mask') callbacks = [ CriterionCallback( input_key="mask", prefix="loss_dice", criterion_key="dice" ), CriterionCallback( input_key="mask", prefix="loss_iou", criterion_key="iou" ), CriterionCallback( input_key="mask", prefix="loss_bce", criterion_key="bce" ), MetricAggregationCallback( prefix="loss", mode="weighted_sum", metrics={"loss_dice": 1.0, "loss_iou": 1.0, "loss_bce": 0.8}, ), # metrics DiceCallback(input_key='mask'), IouCallback(input_key='mask'), ] runner.train( model=model, criterion=criterion, optimizer=optimizer, scheduler=scheduler, loaders=loaders, callbacks=callbacks, logdir=logdir, num_epochs=num_epochs, main_metric="iou", minimize_metric=False, verbose=True, load_best_on_end=True, ) best_model_save_dir = os.path.join(logdir, 'save') os.makedirs(best_model_save_dir, True) torch.save(model, os.path.join(best_model_save_dir, 'best_model.pth')) # save best model (by valid loss) batch = next(iter(loaders["valid"])) try: runner.trace(model=model, batch=batch, logdir=logdir, fp16=False) # optimized version (not all models can be traced) except Exception: pass