magms 1.0.2

PyTorch training manager (v1.0.2)

MAGNET: The First Attempt Towards Modality-Agnostic Networks for Medical Image Segmentation

Currently under review for ISBI 2023

Pre-request

• Python >= 3.9
• torchmanager >= 1.0.4
• Monai >= 0.9.0

Get Started

1. Convert multiple datasets to a magnet.data.TargetDataset and use magnet.data.TargetedDataLoader to load the data

import ...
import magnet

# load data
train_dataset_1, val_dataset_1 = ...
train_dataset_2, val_dataset_2 = ...
...
target_dict = {0: "m1", 1: "m2", ...}
training_dataset = magnet.data.TargetedDataset(train_dataset_1, train_dataset_2, target_dict=target_dict)
training_dataset = magnet.data.TargetedDataLoader(training_dataset, batch_size=batch_size, shuffle=True, num_workers=1, pin_memory=True)
validation_dataset = {
    "m1": data.DataLoader(val_dataset_1, batch_size=batch_size, shuffle=False),
    "m2": data.DataLoader(val_dataset_2, batch_size=batch_size, shuffle=False),
	...
}

2. Simpy build the MAGNET (UNETR backbone) with magnet.build function

num_modalities: int = ...
num_classes: int = ...
img_size: Union[int, Sequence[int]] = ...
model = magnet.build(num_modalities, num_classes, img_size, target_dict=target_dict)

3. Or use the deeper magnet.nn framework to share layers in a torch.nn.Module by given names manually

model1: torch.nn.Module = ...
model2: torch.nn.Module = ...
shared_modules = {
	"layer1": model_to_share.layer1,
	"layer2": model_to_share.layer2,
	...
}
model = magnet.nn.share_modules([model1.some_layers, model2.some_layers], shared_modules, target_dict=target_dict)

4. Compile manager and train/test

optimizer = ...
loss_fn = ...
metric_fns = ...

epochs = ...
callbacks = ...

manager = magnet.Manager(model, optimizer, loss_fn=loss_fn, metric_fns=metric_fns)
manager.fit(training_dataset, epochs, val_dataset=validation_dataset, callbacks=callbacks)
summary.test(validation_dataset)
print(summary)

Monai Support

• Using magnet.MonaigManager instead of Manager
• Post processing support with post_labels and post_predicts

post_labels = [...]
post_predicts = [...]

manager = magnet.MonaigManager(model, post_labels=post_labels, post_predicts=post_predicts, optimizer=optimizer, loss_fn=loss_fn, metric_fns=metric_fns)