Add files via upload

mace-bench/3rdparty/SevenNet/sevenn/train/trainer.py

+import os
+import uuid
+from datetime import datetime
+from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
+
+import torch
+import torch.distributed as dist
+import torch.nn
+from torch.nn.parallel import DistributedDataParallel as DDP
+from tqdm import tqdm
+
+import sevenn._keys as KEY
+from sevenn.error_recorder import ErrorRecorder
+from sevenn.train.loss import LossDefinition
+
+from .loss import get_loss_functions_from_config
+from .optim import optim_dict, scheduler_dict
+
+
+class Trainer:
+    """
+    Training routine specialized for this package. Depends on 'sevenn.train.loss'
+
+    Args:
+        model: model to train
+        loss_functions: List of tuples of [LossDefinition, float]. 'float' is for
+                        loss weight for each Loss function
+        optimizer_cls: torch optimizer class to initialize
+        optimizer_args: optimizer keyword argument except 'param'
+        scheduler_cls: torch scheduler class to initialize, can be None
+        optimizer_args: optimizer keyword argument except 'optimizer'
+        device: device to train model, defaults to 'auto'
+        distributed: whether this is distributed training
+        distributed_backend: torch DDP backend. Should be one of 'nccl', 'mpi'
+    """
+
+    def __init__(
+        self,
+        model: torch.nn.Module,
+        loss_functions: List[Tuple[LossDefinition, float]],
+        optimizer_cls,
+        optimizer_args: Optional[dict] = None,
+        scheduler_cls=None,
+        scheduler_args: Optional[dict] = None,
+        device: Union[torch.device, str] = 'auto',
+        distributed: bool = False,
+        distributed_backend: str = 'nccl',
+    ):
+        if device == 'auto':
+            device = 'cuda' if torch.cuda.is_available() else 'cpu'
+            if distributed_backend == 'mpi':
+                device = 'cpu'
+
+        if distributed:
+            local_rank = int(os.environ['LOCAL_RANK'])
+            self.rank = local_rank
+            if distributed_backend == 'nccl':
+                device = torch.device('cuda', local_rank)
+                self.model = DDP(model.to(device), device_ids=[device])
+            elif distributed_backend == 'mpi':
+                self.model = DDP(model.to(device))
+            else:
+                raise ValueError(f'Unknown DDP backend: {distributed_backend}')
+            dist.barrier()
+            self.model.module.set_is_batch_data(True)
+        else:
+            self.model = model.to(device)
+            self.model.set_is_batch_data(True)
+            self.rank = 0
+
+        self.device = device
+        self.distributed = distributed
+
+        param = [p for p in self.model.parameters() if p.requires_grad]
+        self.optimizer = optimizer_cls(param, **optimizer_args)
+        if scheduler_cls is not None:
+            self.scheduler = scheduler_cls(self.optimizer, **scheduler_args)
+        else:
+            self.scheduler = None
+        self.loss_functions = loss_functions
+
+    @staticmethod
+    def from_config(model: torch.nn.Module, config: Dict[str, Any]) -> 'Trainer':
+        trainer = Trainer(
+            model,
+            loss_functions=get_loss_functions_from_config(config),
+            optimizer_cls=optim_dict[config.get(KEY.OPTIMIZER, 'adam').lower()],
+            optimizer_args=config.get(KEY.OPTIM_PARAM, {}),
+            scheduler_cls=scheduler_dict[
+                config.get(KEY.SCHEDULER, 'exponentiallr').lower()
+            ],
+            scheduler_args=config.get(KEY.SCHEDULER_PARAM, {}),
+            device=config.get(KEY.DEVICE, 'auto'),
+            distributed=config.get(KEY.IS_DDP, False),
+            distributed_backend=config.get(KEY.DDP_BACKEND, 'nccl'),
+        )
+        return trainer
+
+    @staticmethod
+    def args_from_checkpoint(checkpoint: str) -> Tuple[Dict, Dict, Dict]:
+        """
+        Usage:
+            trainer_args, optim_stct, scheduler_stct = args_from_checkpoint('7net-0')
+            # Do what you want to do here
+            trainer = Trainer(**trainer_args)
+            trainer.load_state_dict(
+                optimizer_state_dict=optim_stct,
+                scheduler_state_dict=scheduler_stct,
+        """
+        from sevenn.util import load_checkpoint
+
+        cp = load_checkpoint(checkpoint)
+
+        model = cp.build_model()
+        config = cp.config
+        optimizer_cls = optim_dict[config[KEY.OPTIMIZER].lower()]
+        scheduler_cls = scheduler_dict[config[KEY.SCHEDULER].lower()]
+        loss_functions = get_loss_functions_from_config(config)
+
+        return (
+            {
+                'model': model,
+                'loss_functions': loss_functions,
+                'optimizer_cls': optimizer_cls,
+                'optimizer_args': config[KEY.OPTIM_PARAM],
+                'scheduler_cls': scheduler_cls,
+                'scheduler_args': config[KEY.SCHEDULER_PARAM],
+            },
+            cp.optimizer_state_dict,
+            cp.scheduler_state_dict,
+        )
+
+    def run_one_epoch(
+        self,
+        loader: Iterable,
+        is_train: bool = False,
+        error_recorder: Optional[ErrorRecorder] = None,
+        wrap_tqdm: Union[bool, int] = False,
+    ) -> None:
+        """
+        Run single epoch with given dataloader
+        Args:
+            loader: iterable yieds AtomGraphData
+            is_train: if true, do backward() and optimizer step
+            error_recorder: ErrorRecorder instance to compute errors (RMSEm MAE, ..)
+            wrap_tqdm: wrap given dataloader with tqdm for progress bar
+        """
+        if is_train:
+            self.model.train()
+        else:
+            self.model.eval()
+
+        if wrap_tqdm:
+            total_len = wrap_tqdm if isinstance(wrap_tqdm, int) else None
+            loader = tqdm(loader, total=total_len)
+        for _, batch in enumerate(loader):
+            if is_train:
+                self.optimizer.zero_grad()
+            batch = batch.to(self.device, non_blocking=True)
+            output = self.model(batch)
+            if error_recorder is not None:
+                error_recorder.update(output)
+            if is_train:
+                total_loss = torch.tensor([0.0], device=self.device)
+                for loss_def, w in self.loss_functions:
+                    indv_loss = loss_def.get_loss(output, self.model)
+                    if indv_loss is not None:
+                        total_loss += (indv_loss * w)
+                total_loss.backward()
+                self.optimizer.step()
+
+        if self.distributed and error_recorder is not None:
+            self.recorder_all_reduce(error_recorder)
+
+    def scheduler_step(self, metric: Optional[float] = None) -> None:
+        if self.scheduler is None:
+            return
+        if isinstance(self.scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau):
+            assert isinstance(metric, float)
+            self.scheduler.step(metric)
+        else:
+            self.scheduler.step()
+
+    def get_lr(self) -> float:
+        return float(self.optimizer.param_groups[0]['lr'])
+
+    def recorder_all_reduce(self, recorder: ErrorRecorder) -> None:
+        for metric in recorder.metrics:
+            # metric.value._ddp_reduce(self.device)
+            metric.ddp_reduce(self.device)
+
+    def get_checkpoint_dict(self) -> dict:
+        if self.distributed:
+            model_state_dct = self.model.module.state_dict()
+        else:
+            model_state_dct = self.model.state_dict()
+        return {
+            'model_state_dict': model_state_dct,
+            'optimizer_state_dict': self.optimizer.state_dict(),
+            'scheduler_state_dict': self.scheduler.state_dict()
+            if self.scheduler is not None
+            else None,
+            'time': datetime.now().strftime('%Y-%m-%d %H:%M'),
+            'hash': uuid.uuid4().hex,
+        }
+
+    def write_checkpoint(self, path: str, **extra) -> None:
+        if self.distributed and self.rank != 0:
+            return
+        cp = self.get_checkpoint_dict()
+        cp.update(**extra)
+        torch.save(cp, path)
+
+    def load_state_dicts(
+        self,
+        model_state_dict: Optional[Dict] = None,
+        optimizer_state_dict: Optional[Dict] = None,
+        scheduler_state_dict: Optional[Dict] = None,
+        strict: bool = True,
+    ) -> None:
+        if model_state_dict is not None:
+            if self.distributed:
+                self.model.module.load_state_dict(model_state_dict, strict=strict)
+            else:
+                self.model.load_state_dict(model_state_dict, strict=strict)
+
+        if optimizer_state_dict is not None:
+            self.optimizer.load_state_dict(optimizer_state_dict)
+        if scheduler_state_dict is not None and self.scheduler is not None:
+            self.scheduler.load_state_dict(scheduler_state_dict)