Bring your own components (BYOC)

Clara Train allows researchers to solve new/different problems and innovate by writing their own components in a modular way. In order to do this, users can write their own components in python files then point to these files in the train_config.json file by providing the paths for the new components.

Users can look at the component implementations in MONAI and use any as a template to modify for creating their own components.

For basic components such as networks, loss functions, etc., they do not have to be related to ignite or MONAI workflows as long as the component conforms to the expected PyTorch APIs:

  1. Put the BYOC code into the MMAR/custom folder, for example: my_model.py into a directory named “custom” in the MMAR.

  2. Config your $PYTHONPATH to include the above path (for example, setting export PYTHONPATH="$PYTHONPATH:/opt/nvidia:$MMAR_ROOT/custom" in set_env.sh).

  3. Use your custom component in the MMAR config with “path”: “” in place of “name”: “”.

"model": {
   "path": "my_model.Model",
      "args": {
         ...
   }
}

For other custom logic, we defined 3 levels of APIs to achieve a wide range of customizability:

Handlers: trigger at many Events: started, completed, epoch_started/completed, etc. There are already rich handlers in MONAI: https://docs.monai.io/en/latest/handlers.html

Iteration: write arbitrary computation progress for every iteration.

Extend SupervisedTrainer: write arbitrary trainer to do very complicated logic.

Examples of BYOC

See an example notebook for BYOC.

Bring extra input data for the network

class ExtraInputNet(torch.nn.Module):
   def __init__(self):
       super().__init__()
       self.extra_tensor = None

   def forward(self, image):
       return image + self.extra_tensor


class ExtraInputHandler:
   def attach(self, engine):
       engine.add_event_handler(Events.ITERATION_STARTED, self)

   def __call__(self, engine):
       engine.network.extra_tensor = engine.state.batch["extra"].to(engine.state.device)

Keep init weights of the network for loss

class LossWithWeights(torch.nn.Module):
   def __init__(self):
       super().__init__()
       self.init_weights = None

   def forward(self, pred, label):
       return pred - label + self.init_weights  # fake code


class SaveInitWeightsHandler:
   def attach(self, engine):
       engine.add_event_handler(Events.STARTED, self)

   def __call__(self, engine):
       engine.loss_function.init_weights = copy.deepcopy(engine.network.state_dict())

Add noise to predictions before loss

from monai.engines import IterationEvents

class AddNoiseHandler:
   def attach(self, engine):
       engine.add_event_handler(IterationEvents.FORWARD_COMPLETED, self)

   def __call__(self, engine):
       engine.state.output["pred"] += 1.0

Shutdown program if loss < 0.5

class ShutDownHandler:
   def attach(self, engine):
       engine.add_event_handler(Events.ITERATION_COMPLETED, self)

   def __call__(self, engine):
       if engine.state.output["loss"] < 0.5:
           engine.terminate()

Do optimizer.step() every 10 iterations

from medl.engines.iterations import Iteration

class AccumulateIteration(Iteration):
   def __call__(self, engine, batchdata):
       batch = engine.prepare_batch(batchdata, engine.state.device, engine.non_blocking)
       engine.network.train()
       engine.optimizer.zero_grad()

       pred = engine.inferer(batch[0], engine.network)
       loss = engine.loss_function(pred, batch[1])
       loss.backward()
       if engine.state.iteration % 10 == 0:
           engine.optimizer.step()

       return {"image": batch[0], "label": batch[1], "pred": pred, "loss": loss}

Customize SupervisedTrainer

from ignite.engine import EventEnum
from monai.engines import SupervisedTrainer


class CustomEvents(EventEnum):
   ACCUMULATE_COMPLETED = "accumulate_completed"


class CustomTrainer(SupervisedTrainer):
   def _register_additional_events(self):
       super()._register_additional_events()
       self.register_events(*CustomEvents)

BYOW: BYOC Trainer

If you have needs for custom training logic that is not covered by the training workflow, you can bring your own workflow (BYOW).

For custom training logic, the developer of the MMAR is responsible for what should be in the JSON config and what should be in BYOC python code. For an extreme example, all the logic can be in a BYOC trainer, then this trainer is set in config_train.json:

{
  "epochs": 1260,
  "learning_rate": 2e-4,
  "amp": true,
  "determinism": {
    "random_seed": 0
  },
  "train": {
    "trainer": {
      "path": "byoc_trainer.Trainer",
      "args": {
        "model_path": "{MMAR_CKPT}",
        "amp": "{amp}",
        "num_epochs": "{epochs}",
        "lr": "{learning_rate}"
      }
    }
  }
}