htsn/lung_nodule_ct_detection

1# Copyright (c) MONAI Consortium2# Licensed under the Apache License, Version 2.0 (the "License");3# you may not use this file except in compliance with the License.4# You may obtain a copy of the License at5#     http://www.apache.org/licenses/LICENSE-2.06# Unless required by applicable law or agreed to in writing, software7# distributed under the License is distributed on an "AS IS" BASIS,8# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.9# See the License for the specific language governing permissions and10# limitations under the License.11 12from __future__ import annotations13 14from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, List, Optional, Sequence, Tuple, Union15 16import torch17from monai.engines.trainer import Trainer18from monai.engines.utils import IterationEvents, default_metric_cmp_fn19from monai.inferers import Inferer20from monai.transforms import Transform21from monai.utils import IgniteInfo, min_version, optional_import22from monai.utils.enums import CommonKeys as Keys23from torch.optim.optimizer import Optimizer24from torch.utils.data import DataLoader25 26if TYPE_CHECKING:27    from ignite.engine import Engine, EventEnum28    from ignite.metrics import Metric29else:30    Engine, _ = optional_import("ignite.engine", IgniteInfo.OPT_IMPORT_VERSION, min_version, "Engine")31    Metric, _ = optional_import("ignite.metrics", IgniteInfo.OPT_IMPORT_VERSION, min_version, "Metric")32    EventEnum, _ = optional_import("ignite.engine", IgniteInfo.OPT_IMPORT_VERSION, min_version, "EventEnum")33 34__all__ = ["DetectionTrainer"]35 36 37def detection_prepare_batch(38    batchdata: List[Dict[str, torch.Tensor]],39    device: Optional[Union[str, torch.device]] = None,40    non_blocking: bool = False,41    **kwargs,42) -> Union[Tuple[torch.Tensor, Optional[torch.Tensor]], torch.Tensor]:43    """44    Default function to prepare the data for current iteration.45    Args `batchdata`, `device`, `non_blocking` refer to the ignite API:46    https://pytorch.org/ignite/v0.4.8/generated/ignite.engine.create_supervised_trainer.html.47    `kwargs` supports other args for `Tensor.to()` API.48    Returns:49        image, label(optional).50    """51    inputs = [52        batch_data_ii["image"].to(device=device, non_blocking=non_blocking, **kwargs)53        for batch_data_i in batchdata54        for batch_data_ii in batch_data_i55    ]56 57    if isinstance(batchdata[0][0].get(Keys.LABEL), torch.Tensor):58        targets = [59            dict(60                label=batch_data_ii["label"].to(device=device, non_blocking=non_blocking, **kwargs),61                box=batch_data_ii["box"].to(device=device, non_blocking=non_blocking, **kwargs),62            )63            for batch_data_i in batchdata64            for batch_data_ii in batch_data_i65        ]66        return (inputs, targets)67    return inputs, None68 69 70class DetectionTrainer(Trainer):71    """72    Supervised detection training method with image and label, inherits from ``Trainer`` and ``Workflow``.73    Args:74        device: an object representing the device on which to run.75        max_epochs: the total epoch number for trainer to run.76        train_data_loader: Ignite engine use data_loader to run, must be Iterable or torch.DataLoader.77        detector: detector to train in the trainer, should be regular PyTorch `torch.nn.Module`.78        optimizer: the optimizer associated to the detector, should be regular PyTorch optimizer from `torch.optim`79            or its subclass.80        epoch_length: number of iterations for one epoch, default to `len(train_data_loader)`.81        non_blocking: if True and this copy is between CPU and GPU, the copy may occur asynchronously82            with respect to the host. For other cases, this argument has no effect.83        prepare_batch: function to parse expected data (usually `image`,`box`, `label` and other detector args)84            from `engine.state.batch` for every iteration, for more details please refer to:85            https://pytorch.org/ignite/generated/ignite.engine.create_supervised_trainer.html.86        iteration_update: the callable function for every iteration, expect to accept `engine`87            and `engine.state.batch` as inputs, return data will be stored in `engine.state.output`.88            if not provided, use `self._iteration()` instead. for more details please refer to:89            https://pytorch.org/ignite/generated/ignite.engine.engine.Engine.html.90        inferer: inference method that execute model forward on input data, like: SlidingWindow, etc.91        postprocessing: execute additional transformation for the model output data.92            Typically, several Tensor based transforms composed by `Compose`.93        key_train_metric: compute metric when every iteration completed, and save average value to94            engine.state.metrics when epoch completed. key_train_metric is the main metric to compare and save the95            checkpoint into files.96        additional_metrics: more Ignite metrics that also attach to Ignite Engine.97        metric_cmp_fn: function to compare current key metric with previous best key metric value,98            it must accept 2 args (current_metric, previous_best) and return a bool result: if `True`, will update99            `best_metric` and `best_metric_epoch` with current metric and epoch, default to `greater than`.100        train_handlers: every handler is a set of Ignite Event-Handlers, must have `attach` function, like:101            CheckpointHandler, StatsHandler, etc.102        amp: whether to enable auto-mixed-precision training, default is False.103        event_names: additional custom ignite events that will register to the engine.104            new events can be a list of str or `ignite.engine.events.EventEnum`.105        event_to_attr: a dictionary to map an event to a state attribute, then add to `engine.state`.106            for more details, check: https://pytorch.org/ignite/generated/ignite.engine.engine.Engine.html107            #ignite.engine.engine.Engine.register_events.108        decollate: whether to decollate the batch-first data to a list of data after model computation,109            recommend `decollate=True` when `postprocessing` uses components from `monai.transforms`.110            default to `True`.111        optim_set_to_none: when calling `optimizer.zero_grad()`, instead of setting to zero, set the grads to None.112            more details: https://pytorch.org/docs/stable/generated/torch.optim.Optimizer.zero_grad.html.113        to_kwargs: dict of other args for `prepare_batch` API when converting the input data, except for114            `device`, `non_blocking`.115        amp_kwargs: dict of the args for `torch.cuda.amp.autocast()` API, for more details:116            https://pytorch.org/docs/stable/amp.html#torch.cuda.amp.autocast.117    """118 119    def __init__(120        self,121        device: torch.device,122        max_epochs: int,123        train_data_loader: Iterable | DataLoader,124        detector: torch.nn.Module,125        optimizer: Optimizer,126        epoch_length: int | None = None,127        non_blocking: bool = False,128        prepare_batch: Callable = detection_prepare_batch,129        iteration_update: Callable[[Engine, Any], Any] | None = None,130        inferer: Inferer | None = None,131        postprocessing: Transform | None = None,132        key_train_metric: dict[str, Metric] | None = None,133        additional_metrics: dict[str, Metric] | None = None,134        metric_cmp_fn: Callable = default_metric_cmp_fn,135        train_handlers: Sequence | None = None,136        amp: bool = False,137        event_names: list[str | EventEnum] | None = None,138        event_to_attr: dict | None = None,139        decollate: bool = True,140        optim_set_to_none: bool = False,141        to_kwargs: dict | None = None,142        amp_kwargs: dict | None = None,143    ) -> None:144        super().__init__(145            device=device,146            max_epochs=max_epochs,147            data_loader=train_data_loader,148            epoch_length=epoch_length,149            non_blocking=non_blocking,150            prepare_batch=prepare_batch,151            iteration_update=iteration_update,152            postprocessing=postprocessing,153            key_metric=key_train_metric,154            additional_metrics=additional_metrics,155            metric_cmp_fn=metric_cmp_fn,156            handlers=train_handlers,157            amp=amp,158            event_names=event_names,159            event_to_attr=event_to_attr,160            decollate=decollate,161            to_kwargs=to_kwargs,162            amp_kwargs=amp_kwargs,163        )164 165        self.detector = detector166        self.optimizer = optimizer167        self.optim_set_to_none = optim_set_to_none168 169    def _iteration(self, engine, batchdata: dict[str, torch.Tensor]):170        """171        Callback function for the Supervised Training processing logic of 1 iteration in Ignite Engine.172        Return below items in a dictionary:173            - IMAGE: image Tensor data for model input, already moved to device.174            - BOX: box regression loss corresponding to the image, already moved to device.175            - LABEL: classification loss corresponding to the image, already moved to device.176            - LOSS: weighted sum of loss values computed by loss function.177        Args:178            engine: `DetectionTrainer` to execute operation for an iteration.179            batchdata: input data for this iteration, usually can be dictionary or tuple of Tensor data.180        Raises:181            ValueError: When ``batchdata`` is None.182        """183 184        if batchdata is None:185            raise ValueError("Must provide batch data for current iteration.")186 187        batch = engine.prepare_batch(batchdata, engine.state.device, engine.non_blocking, **engine.to_kwargs)188        if len(batch) == 2:189            inputs, targets = batch190            args: tuple = ()191            kwargs: dict = {}192        else:193            inputs, targets, args, kwargs = batch194        # put iteration outputs into engine.state195        engine.state.output = {Keys.IMAGE: inputs, Keys.LABEL: targets}196 197        def _compute_pred_loss(w_cls: float = 1.0, w_box_reg: float = 1.0):198            """199            Args:200                w_cls: weight of classification loss201                w_box_reg: weight of box regression loss202            """203            outputs = engine.detector(inputs, targets)204            engine.state.output[engine.detector.cls_key] = outputs[engine.detector.cls_key]205            engine.state.output[engine.detector.box_reg_key] = outputs[engine.detector.box_reg_key]206            engine.state.output[Keys.LOSS] = (207                w_cls * outputs[engine.detector.cls_key] + w_box_reg * outputs[engine.detector.box_reg_key]208            )209            engine.fire_event(IterationEvents.LOSS_COMPLETED)210 211        engine.detector.train()212        engine.optimizer.zero_grad(set_to_none=engine.optim_set_to_none)213 214        if engine.amp and engine.scaler is not None:215            with torch.cuda.amp.autocast(**engine.amp_kwargs):216                inputs = [img.to(torch.float16) for img in inputs]217                _compute_pred_loss()218            engine.scaler.scale(engine.state.output[Keys.LOSS]).backward()219            engine.fire_event(IterationEvents.BACKWARD_COMPLETED)220            engine.scaler.step(engine.optimizer)221            engine.scaler.update()222        else:223            _compute_pred_loss()224            engine.state.output[Keys.LOSS].backward()225            engine.fire_event(IterationEvents.BACKWARD_COMPLETED)226            engine.optimizer.step()227 228        return engine.state.output229