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
