GraFEIIgniteTrainer Class Reference

Public Member Functions
def	__init__ (self, model, optimizer, loss_fn, device, configs, tags, scheduler=None, ignore_index=-1.0)
def	setup_handlers (self, cfg_filename="config.yaml")
def	log_results (self, trainer, mode_tags)

Public Attributes
	model
	Model.
	optimizer
	Optimizer.
	configs
	Configs.
	tags
	Tags.
	ignore_index
	Index to ignore.
	device
	CPU or GPU.
	timestamp
	Run timestamp to distinguish trainings.
	run_dir
	Output directory for checkpoints.
	trainer
	Ignite trainer.
	evaluators
	Setup train and validation evaluators.

Protected Member Functions
def	_score_fn (self, engine)
def	_perfect_score_fn (self, engine)
def	_clean_config_dict (self, configs)

Detailed Description

Class to setup the ignite trainer and hold all the things associated.

:param model: The actual PyTorch model.
:type model: `Model <https://pytorch.org/tutorials/beginner/introyt/modelsyt_tutorial.html>`_
:param optimizer: Optimizer used in training.
:type optimizer: `Optimizer <https://pytorch.org/docs/stable/optim.html#torch.optim.Optimizer>`_
:param loss_fn: Loss function.
:type loss_fn: `Loss <https://pytorch.org/docs/stable/nn.html#loss-functions>`_
:param device: Device to use.
:type device: `Device <https://pytorch.org/docs/stable/tensor_attributes.html#torch.device>`_
:param configs: Dictionary of run configs from loaded yaml config file.
:type configs: dict
:param tags: Various tags to sort train and validation evaluators by, e.g. "Training", "Validation".
:type tags: list
:param scheduler: Learning rate scheduler.
:type scheduler: `Scheduler <https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate>`_
:param ignore_index: Label index to ignore when calculating metrics, e.g. padding.
:type ignore_index: int

Definition at line 21 of file create_trainer.py.

Constructor & Destructor Documentation

__init__()

def __init__	(		self,
			model,
			optimizer,
			loss_fn,
			device,
			configs,
			tags,
			scheduler = None,
			ignore_index = -1.0
	)

Initialization.

Definition at line 43 of file create_trainer.py.

    ):
        """
        Initialization.
        """
        
        self.model = model
        
        self.optimizer = optimizer
        
        self.configs = configs
        
        self.tags = tags
        
        self.ignore_index = ignore_index
        
        self.device = device
 
        
        self.timestamp = datetime.now().strftime("%Y.%m.%d_%H.%M")
 
        
        self.run_dir = None
        if self.configs["output"] is not None:
            if ("path" in self.configs["output"].keys()) and (
                self.configs["output"]["path"] is not None
            ):
                self.run_dir = Path(
                    self.configs["output"]["path"],
                    self.configs["output"]["run_name"],
                )
 
        # Setup ignite trainer
        use_amp = configs["train"]["mixed_precision"] and self.device == torch.device(
            "cuda"
        )
 
        if use_amp:
            from torch.cuda.amp import autocast
            from torch.cuda.amp import GradScaler
 
            scaler = GradScaler(enabled=True)
 
        def _update_model(engine, batch):
            # This just sets the training mode
            model.train()
 
            optimizer.zero_grad()
 
            batch = (
                Batch.from_data_list(batch).to(device)
                if isinstance(batch, list)
                else batch.to(device)
            )
 
            x_y, edge_y, u_y = batch.x_y, batch.edge_y, batch.u_y
 
            if use_amp:
                with autocast(enabled=True):
                    x_pred, e_pred, u_pred = model(batch)
                    loss = loss_fn(x_pred, x_y, e_pred, edge_y, u_pred, u_y)
                scaler.scale(loss).backward()
                scaler.step(optimizer)
                scaler.update()
            else:
                x_pred, e_pred, u_pred = model(batch)
                loss = loss_fn(x_pred, x_y, e_pred, edge_y, u_pred, u_y)
                loss.backward()
                optimizer.step()
 
            return loss.item()
 
        
        self.trainer = ignite.engine.Engine(_update_model)
 
        if scheduler:
            ig_scheduler = ignite.handlers.param_scheduler.LRScheduler(scheduler)
            self.trainer.add_event_handler(ignite.engine.Events.ITERATION_STARTED, ig_scheduler)
 
        
        self.evaluators = {}
 
        for tag in self.tags:
            # Setup metrics
            metrics = {
                # ignite.metrics.Loss takes (y_pred, y, **kwargs) arguments.
                # MultiTrainLoss needs in total 6 arguments to be computed,
                # so the additional ones are passed in a dictionary.
                "loss": ignite.metrics.Loss(
                    loss_fn,
                    output_transform=lambda x: [
                        x[0], x[3],
                        {
                            "edge_input": x[1],
                            "edge_target": x[4],
                            "u_input": x[2],
                            "u_target": x[5],
                        },
                    ],
                    device=device,
                ),
                "perfectLCA": PerfectLCA(
                    ignore_index=ignore_index,
                    device=device,
                    output_transform=lambda x: [
                        x[1], x[4], x[6], x[5], x[7], x[8],
                    ],
                ),
                "perfectMasses": PerfectMasses(
                    ignore_index=ignore_index,
                    device=device,
                    output_transform=lambda x: [x[0], x[3], x[5], x[7], x[8]],
                ),
                "perfectEvent": PerfectEvent(
                    ignore_index=ignore_index,
                    device=device,
                    output_transform=lambda x: [
                        x[0], x[3], x[1], x[4], x[6], x[5], x[7], x[8],
                    ],
                ),
            }
 
            def _predict_on_batch(engine, batch):
                model.eval()  # It just enables evaluation mode
 
                batch = (
                    Batch.from_data_list(batch).to(device)
                    if isinstance(batch, list)
                    else batch.to(device)
                )
 
                x_y, edge_y, u_y, edge_index, torch_batch = (
                    batch.x_y,
                    batch.edge_y,
                    batch.u_y,
                    batch.edge_index,
                    batch.batch,
                )
                num_graph = batch.batch[torch_batch.shape[0] - 1] + 1
 
                with torch.no_grad():
                    if use_amp:
                        with autocast(enabled=True):
                            x_pred, e_pred, u_pred = model(batch)
                    else:
                        x_pred, e_pred, u_pred = model(batch)
 
                return (
                    x_pred,
                    e_pred,
                    u_pred,
                    x_y,
                    edge_y,
                    u_y,
                    edge_index,
                    torch_batch,
                    num_graph,
                )
 
            self.evaluators[tag] = ignite.engine.Engine(_predict_on_batch)
 
            for metric_name, metric in zip(metrics.keys(), metrics.values()):
                metric.attach(self.evaluators[tag], metric_name)
 

Member Function Documentation

_clean_config_dict()

def _clean_config_dict (   self,   configs  )

protected

Clean configs to prepare them for writing to file.

Definition at line 224 of file create_trainer.py.

    def _clean_config_dict(self, configs):
        """
        Clean configs to prepare them for writing to file.
        """
        for k, v in configs.items():
            if isinstance(v, collections.abc.Mapping):
                configs[k] = self._clean_config_dict(configs[k])
            elif isinstance(v, np.ndarray):
                configs[k] = v.tolist()
            else:
                configs[k] = v
        return configs
 

_perfect_score_fn()

def _perfect_score_fn (   self,   engine  )

protected

Metric to use for checkpoints

Definition at line 220 of file create_trainer.py.

    def _perfect_score_fn(self, engine):
        """Metric to use for checkpoints"""
        return engine.state.metrics["perfectEvent"]
 

_score_fn()

def _score_fn (   self,   engine  )

protected

Metric to use for early stoppging

Definition at line 216 of file create_trainer.py.

    def _score_fn(self, engine):
        """Metric to use for early stoppging"""
        return -engine.state.metrics["loss"]
 

log_results()

def log_results	(		self,
			trainer,
			mode_tags
	)

Callback to run evaluation and report the results.

:param trainer: Trainer passed by ignite to this method.
:type trainer: `Engine <https://pytorch.org/ignite/generated/ignite.engine.engine.Engine.html#ignite.engine.engine.Engine>`_
:param mode_tags: Dictionary of mode tags containing (mode, dataset, dataloader) tuples.
:type mode_tags: dict

Definition at line 296 of file create_trainer.py.

    def log_results(self, trainer, mode_tags):
        """
        Callback to run evaluation and report the results.
 
        :param trainer: Trainer passed by ignite to this method.
        :type trainer: `Engine <https://pytorch.org/ignite/generated/ignite.engine.engine.Engine.html#ignite.engine.engine.Engine>`_
        :param mode_tags: Dictionary of mode tags containing (mode, dataset, dataloader) tuples.
        :type mode_tags: dict
        """
 
        for tag, values in mode_tags.items():
            evaluator = self.evaluators[tag]
 
            # Need to wrap this in autocast since it calculates metrics (i.e. loss) without autocast switched on
            # This is mostly fine except it fails to correctly cast the class weights tensor passed to the loss
            if self.configs["train"]["mixed_precision"] and self.device == torch.device("cuda"):
                with torch.cuda.amp.autocast():
                    evaluator.run(values[2], epoch_length=None)
            else:
                evaluator.run(values[2], epoch_length=None)
 
            metrics = evaluator.state.metrics
            message = [f"{tag} Results - Epoch: {trainer.state.epoch}"]
            message.extend([f"Avg {m}: {metrics[m]:.4f}" for m in metrics])
            print(message)

setup_handlers()

def setup_handlers	(		self,
			cfg_filename = "config.yaml"
	)

Creates the various ignite handlers (callbacks).

Args:
    cfg_filename (str): Name of config yaml file to use when saving configs.

Definition at line 237 of file create_trainer.py.

    def setup_handlers(self, cfg_filename="config.yaml"):
        """
        Creates the various ignite handlers (callbacks).
 
        Args:
            cfg_filename (str): Name of config yaml file to use when saving configs.
        """
        # Create the output directory
        if self.run_dir is not None:
            self.run_dir.mkdir(parents=True, exist_ok=True)
            # And save the configs, putting here to only save when setting up checkpointing
            with open(
                self.run_dir / f"{self.timestamp}_{cfg_filename}", "w"
            ) as outfile:
                cleaned_configs = self._clean_config_dict(self.configs)
                yaml.dump(cleaned_configs, outfile, default_flow_style=False)
 
        # Setup early stopping
        early_handler = ignite.handlers.EarlyStopping(
            patience=self.configs["train"]["early_stop_patience"],
            score_function=self._score_fn,
            trainer=self.trainer,
            min_delta=1e-3,
        )
        self.evaluators["Validation"].add_event_handler(
            ignite.engine.Events.EPOCH_COMPLETED, early_handler
        )
 
        # Configure saving the best performing model
        if self.run_dir is not None:
            to_save = {
                "model": self.model,
                "optimizer": self.optimizer,
                "trainer": self.trainer,
            }
            # Note that we judge early stopping above by the validation loss, but save the best model
            # according to validation perfectEvent score. This lets training continue for perfectEvent plateaus
            # so long as the model is still changing (and hopefully improving again after some time).
            best_model_handler = ignite.handlers.Checkpoint(
                to_save=to_save,
                save_handler=ignite.handlers.DiskSaver(
                    self.run_dir, create_dir=True, require_empty=False
                ),
                filename_prefix=self.timestamp,
                score_function=self._perfect_score_fn,
                score_name="validation_perfectEvent",
                n_saved=1,
                global_step_transform=ignite.handlers.global_step_from_engine(
                    self.evaluators["Validation"]
                ),
            )
            self.evaluators["Validation"].add_event_handler(
                ignite.engine.Events.EPOCH_COMPLETED, best_model_handler
            )
 
        return
 

Member Data Documentation

configs

configs

Configs.

Definition at line 62 of file create_trainer.py.

device

device

CPU or GPU.

Definition at line 68 of file create_trainer.py.

evaluators

evaluators

Setup train and validation evaluators.

Definition at line 132 of file create_trainer.py.

ignore_index

ignore_index

Index to ignore.

Definition at line 66 of file create_trainer.py.

model

model

Model.

Definition at line 58 of file create_trainer.py.

optimizer

optimizer

Optimizer.

Definition at line 60 of file create_trainer.py.

run_dir

run_dir

Output directory for checkpoints.

Definition at line 74 of file create_trainer.py.

tags

tags

Tags.

Definition at line 64 of file create_trainer.py.

timestamp

timestamp

Run timestamp to distinguish trainings.

Definition at line 71 of file create_trainer.py.

trainer

trainer

Ignite trainer.

Definition at line 125 of file create_trainer.py.

---

The documentation for this class was generated from the following file:

• analysis/scripts/grafei/model/create_trainer.py