diff --git a/CHANGELOG.md b/CHANGELOG.md index 5c6f73cb..14fabea4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,11 @@ All notable changes to this project will be documented in this file. ## [0.13.1.dev] +### Added +- Training features used for SevenNet-Omni: batch training, `OrderedSampler`, `grad_clip`, `onecyclelr` +- Loss: MAE, L2MAE +- Dataset type: `aselmdb`, `custom` + ### Fixed - `D3Calculator()` segfault bug when reusing the calculator within different sized `Atoms`. @@ -39,13 +44,12 @@ All notable changes to this project will be documented in this file. ### Added - TorchSim interface and docs +- SevenNet-Omni-i8, SevenNet-Omni-i12 ## [0.12.0] ### Added - Documentation moved to RTD - LAMMPS-MLIAP integration with GhostExchangeOp - -### Added - SevenNet-Omni - Example config for fine-tuning the SevenNet-MF-ompa model - FlashTP support (https://github.com/SNU-ARC/flashTP) diff --git a/pyproject.toml b/pyproject.toml index 416d4a35..daeea096 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,7 +30,9 @@ dependencies = [ "pandas", "requests", "ninja", - "setuptools>=61.0" + "setuptools>=61.0", + "lmdb<2.0.0", + "orjson" ] [project.optional-dependencies] test = ["pytest", "pytest-cov>=5", "ipython"] diff --git a/setup.cfg b/setup.cfg index 1505c8b8..dc23e4c3 100644 --- a/setup.cfg +++ b/setup.cfg @@ -10,5 +10,5 @@ include_trailing_comma=True force_grid_wrap=0 use_parentheses=True line_length=80 -known_third_party=ase,braceexpand,e3nn,numpy,packaging,pandas,pytest,requests,sklearn,torch,torch_geometric,tqdm,yaml +known_third_party=ase,braceexpand,e3nn,lmdb,numpy,orjson,packaging,pandas,pytest,requests,sklearn,torch,torch_geometric,tqdm,yaml known_first_party= diff --git a/sevenn/_const.py b/sevenn/_const.py index 0e411d77..6c78f9c8 100644 --- a/sevenn/_const.py +++ b/sevenn/_const.py @@ -15,6 +15,13 @@ IMPLEMENTED_SELF_CONNECTION_TYPE = ['nequip', 'linear'] IMPLEMENTED_INTERACTION_TYPE = ['nequip'] +IMPLEMENTED_MODAL_MODULE_DICT = { + KEY.USE_MODAL_NODE_EMBEDDING: 'onehot_to_feature_x', + KEY.USE_MODAL_SELF_INTER_INTRO: 'self_interaction_1', + KEY.USE_MODAL_SELF_INTER_OUTRO: 'self_interaction_2', + KEY.USE_MODAL_OUTPUT_BLOCK: 'reduce_input_to_hidden', +} + IMPLEMENTED_SHIFT = ['per_atom_energy_mean', 'elemwise_reference_energies'] IMPLEMENTED_SCALE = ['force_rms', 'per_atom_energy_std', 'elemwise_force_rms'] @@ -26,6 +33,8 @@ 'Stress', 'Stress_GPa', 'TotalLoss', + 'L2_modal', + 'Modal_cos', ] IMPLEMENTED_MODEL = ['E3_equivariant_model'] @@ -116,6 +125,8 @@ def error_record_condition(x): KEY.CONV_DENOMINATOR: 'avg_num_neigh', KEY.TRAIN_DENOMINTAOR: False, KEY.TRAIN_SHIFT_SCALE: False, + KEY.TRAIN_SHIFT: False, + KEY.TRAIN_SCALE: False, # KEY.OPTIMIZE_BY_REDUCE: True, # deprecated, always True KEY.USE_BIAS_IN_LINEAR: False, KEY.USE_MODAL_NODE_EMBEDDING: False, @@ -158,6 +169,8 @@ def error_record_condition(x): ], KEY.CONVOLUTION_WEIGHT_NN_HIDDEN_NEURONS: list, KEY.TRAIN_SHIFT_SCALE: bool, + KEY.TRAIN_SHIFT: bool, + KEY.TRAIN_SCALE: bool, KEY.TRAIN_DENOMINTAOR: bool, KEY.USE_BIAS_IN_LINEAR: bool, KEY.USE_MODAL_NODE_EMBEDDING: bool, @@ -215,6 +228,7 @@ def model_defaults(config): KEY.REHEARSAL: False, KEY.MEM_BATCH_SIZE: 0, KEY.MEM_RATIO: 1, + KEY.LOADER_KWARGS: {}, # KEY.DATA_SHUFFLE: True, # KEY.DATA_WEIGHT: False, # KEY.DATA_MODALITY: False, @@ -230,7 +244,7 @@ def model_defaults(config): KEY.RATIO: float, KEY.BATCH_SIZE: int, KEY.PREPROCESS_NUM_CORES: int, - KEY.DATASET_TYPE: lambda x: x in ['graph', 'atoms'], + KEY.DATASET_TYPE: lambda x: x in ['graph', 'atoms', 'custom', 'aselmdb'], # KEY.USE_SPECIES_WISE_SHIFT_SCALE: bool, KEY.SHIFT: lambda x: type(x) in [float, list] or x in IMPLEMENTED_SHIFT, KEY.SCALE: lambda x: type(x) in [float, list] or x in IMPLEMENTED_SCALE, @@ -264,15 +278,20 @@ def data_defaults(config): KEY.OPTIM_PARAM: {}, KEY.SCHEDULER: 'exponentiallr', KEY.SCHEDULER_PARAM: {}, + KEY.ENERGY_WEIGHT: 1.0, KEY.FORCE_WEIGHT: 0.1, KEY.STRESS_WEIGHT: 1e-6, # SIMPLE-NN default + KEY.GRAD_CLIP: None, + KEY.REG_PARAM: {}, KEY.PER_EPOCH: 5, + KEY.TRAIN_BY_BATCH: False, # KEY.USE_TESTSET: False, KEY.CONTINUE: { KEY.CHECKPOINT: False, KEY.RESET_OPTIMIZER: False, KEY.RESET_SCHEDULER: False, KEY.RESET_EPOCH: False, + KEY.RESET_DATA_PROGRESS: True, KEY.USE_STATISTIC_VALUES_OF_CHECKPOINT: True, KEY.USE_STATISTIC_VALUES_FOR_CP_MODAL_ONLY: True, }, @@ -296,16 +315,21 @@ def data_defaults(config): TRAINING_CONFIG_CONDITION = { KEY.RANDOM_SEED: int, KEY.EPOCH: int, + KEY.ENERGY_WEIGHT: float, KEY.FORCE_WEIGHT: float, KEY.STRESS_WEIGHT: float, + KEY.GRAD_CLIP: lambda x: x is None or (type(x) in [float, int] and x > 0), + KEY.REG_PARAM: dict, KEY.USE_TESTSET: None, # Not used KEY.NUM_WORKERS: int, - KEY.PER_EPOCH: int, + KEY.PER_EPOCH: lambda x: type(x) in [float, int], + KEY.TRAIN_BY_BATCH: bool, KEY.CONTINUE: { KEY.CHECKPOINT: str, KEY.RESET_OPTIMIZER: bool, KEY.RESET_SCHEDULER: bool, KEY.RESET_EPOCH: bool, + KEY.RESET_DATA_PROGRESS: bool, KEY.USE_STATISTIC_VALUES_OF_CHECKPOINT: bool, KEY.USE_STATISTIC_VALUES_FOR_CP_MODAL_ONLY: bool, }, diff --git a/sevenn/_keys.py b/sevenn/_keys.py index 31d29ae5..8773a099 100644 --- a/sevenn/_keys.py +++ b/sevenn/_keys.py @@ -117,12 +117,17 @@ EPOCH = 'epoch' LOSS = 'loss' LOSS_PARAM = 'loss_param' +LOSS_TYPE = 'loss_type' +LOSS_WEIGHT = 'loss_weight' OPTIMIZER = 'optimizer' OPTIM_PARAM = 'optim_param' SCHEDULER = 'scheduler' SCHEDULER_PARAM = 'scheduler_param' +SCHEDULER_BATCH_MODE = 'scheduler_batch_mode' +ENERGY_WEIGHT = 'energy_loss_weight' FORCE_WEIGHT = 'force_loss_weight' STRESS_WEIGHT = 'stress_loss_weight' +GRAD_CLIP = 'grad_clip' DEVICE = 'device' DTYPE = 'dtype' @@ -135,6 +140,7 @@ RESET_OPTIMIZER = 'reset_optimizer' RESET_SCHEDULER = 'reset_scheduler' RESET_EPOCH = 'reset_epoch' +RESET_DATA_PROGRESS = 'reset_data_progress' USE_STATISTIC_VALUES_OF_CHECKPOINT = 'use_statistic_values_of_checkpoint' USE_STATISTIC_VALUES_FOR_CP_MODAL_ONLY = ( 'use_statistic_values_for_cp_modal_only' @@ -157,6 +163,11 @@ DDP_BACKEND = 'ddp_backend' PER_EPOCH = 'per_epoch' +TRAIN_BY_BATCH = 'train_by_batch' +TOTAL_DATA_NUM = 'total_data_num' +CURRENT_DATA_IDX = 'current_data_index' +NUMPY_RNG_STATE = 'numpy_rng_state' + USE_WEIGHT = 'use_weight' USE_MODALITY = 'use_modality' DEFAULT_MODAL = 'default_modal' @@ -218,12 +229,15 @@ CONV_DENOMINATOR = 'conv_denominator' SHIFT = 'shift' SCALE = 'scale' +LOADER_KWARGS = 'loader_kwargs' USE_SPECIES_WISE_SHIFT_SCALE = 'use_species_wise_shift_scale' USE_MODAL_WISE_SHIFT = 'use_modal_wise_shift' USE_MODAL_WISE_SCALE = 'use_modal_wise_scale' TRAIN_SHIFT_SCALE = 'train_shift_scale' +TRAIN_SHIFT = 'train_shift' +TRAIN_SCALE = 'train_scale' TRAIN_DENOMINTAOR = 'train_denominator' INTERACTION_TYPE = 'interaction_type' TRAIN_AVG_NUM_NEIGH = 'train_avg_num_neigh' # deprecated @@ -232,6 +246,9 @@ CUEQUIVARIANCE_CONFIG = 'cuequivariance_config' USE_OEQ = 'use_oeq' +REG_PARAM = 'regularization_param' +REG_WEIGHT = 'regularization_weight' + _NORMALIZE_SPH = '_normalize_sph' OPTIMIZE_BY_REDUCE = 'optimize_by_reduce' diff --git a/sevenn/checkpoint.py b/sevenn/checkpoint.py index 0d0d57de..8e4a704e 100644 --- a/sevenn/checkpoint.py +++ b/sevenn/checkpoint.py @@ -226,6 +226,7 @@ def __init__(self, checkpoint_path: Union[pathlib.Path, str]) -> None: self._checkpoint_path = os.path.abspath(checkpoint_path) self._config = None self._epoch = None + self._data_progress = None self._model_state_dict = None self._optimizer_state_dict = None self._scheduler_state_dict = None @@ -304,6 +305,12 @@ def epoch(self) -> Optional[int]: self._load() return self._epoch + @property + def data_progress(self) -> Optional[Dict[str, int]]: + if not self._loaded: + self._load() + return self._data_progress + @property def time(self) -> str: if not self._loaded: @@ -328,6 +335,7 @@ def _load(self) -> None: self._optimizer_state_dict = cp.get('optimizer_state_dict', {}) self._scheduler_state_dict = cp.get('scheduler_state_dict', {}) self._epoch = cp.get('epoch', None) + self._data_progress = cp.get('data_progress', None) self._time = cp.get('time', 'Not found') self._hash = cp.get('hash', 'Not found') diff --git a/sevenn/error_recorder.py b/sevenn/error_recorder.py index 262ea06f..8ebf4d8b 100644 --- a/sevenn/error_recorder.py +++ b/sevenn/error_recorder.py @@ -60,6 +60,18 @@ 'coeff': 160.21766208, 'vdim': 6, }, + 'L2_modal': { + 'name': 'L2_modal', + 'ref_key': None, + 'pred_key': None, + 'unit': None, + }, + 'Modal_cos': { + 'name': 'Modal_cos', + 'ref_key': None, + 'pred_key': None, + 'unit': None, + }, 'TotalLoss': { 'name': 'TotalLoss', 'unit': None, @@ -83,9 +95,15 @@ def __init__(self): self._sum = 0.0 self._count = 0 - def update(self, values: torch.Tensor) -> None: - self._sum += values.sum().item() - self._count += values.numel() + def update(self, values: Union[torch.Tensor, float]) -> None: + if isinstance(values, torch.Tensor): + self._sum += values.sum().item() + self._count += values.numel() + elif isinstance(values, float): + self._sum += values + self._count += 1 + else: + raise ValueError(f'Unsupported type: {type(values)}') def _ddp_reduce(self, device): _sum = torch.tensor(self._sum, device=device) @@ -103,8 +121,9 @@ def get(self) -> float: class ErrorMetric: """ - Base class for error metrics We always average error by # of structures, - and designed to collect errors in the middle of iteration (by AverageNumber) + Base class for error metrics + Always average error by # of structures, + Designed to collect errors in the middle of iteration (by AverageNumber) """ def __init__( @@ -127,7 +146,9 @@ def __init__( self.ignore_unlabeled = ignore_unlabeled self.value = AverageNumber() - def update(self, output: 'AtomGraphData') -> None: + def update( + self, output: 'AtomGraphData', model: Optional[Callable] = None + ) -> None: raise NotImplementedError def _retrieve( @@ -180,7 +201,9 @@ def _square_error( ) -> torch.Tensor: return self._se(y_ref.view(-1, vdim), y_pred.view(-1, vdim)).sum(dim=1) - def update(self, output: 'AtomGraphData') -> None: + def update( + self, output: 'AtomGraphData', model: Optional[Callable] = None + ) -> None: y_ref, y_pred = self._retrieve(output) se = self._square_error(y_ref, y_pred, self.vdim) self.value.update(se) @@ -204,7 +227,9 @@ def _square_error( ) -> torch.Tensor: return self._se(y_ref, y_pred) - def update(self, output: 'AtomGraphData') -> None: + def update( + self, output: 'AtomGraphData', model: Optional[Callable] = None + ) -> None: y_ref, y_pred = self._retrieve(output) y_ref = y_ref.view(-1) y_pred = y_pred.view(-1) @@ -228,7 +253,9 @@ def _square_error( ) -> torch.Tensor: return torch.abs(y_ref - y_pred) - def update(self, output: 'AtomGraphData') -> None: + def update( + self, output: 'AtomGraphData', model: Optional[Callable] = None + ) -> None: y_ref, y_pred = self._retrieve(output) y_ref = y_ref.reshape((-1,)) y_pred = y_pred.reshape((-1,)) @@ -248,7 +275,9 @@ def __init__(self, func: Callable, **kwargs) -> None: super().__init__(**kwargs) self.func = func - def update(self, output: 'AtomGraphData') -> None: + def update( + self, output: 'AtomGraphData', model: Optional[Callable] = None + ) -> None: y_ref, y_pred = self._retrieve(output) se = self.func(y_ref, y_pred) if len(y_ref) > 0 else torch.tensor([]) self.value.update(se) @@ -272,11 +301,51 @@ def __init__( ) self.loss_def = loss_def - def update(self, output: 'AtomGraphData') -> None: - loss = self.loss_def.get_loss(output) # type: ignore + def update( + self, output: 'AtomGraphData', model: Optional[Callable] = None + ) -> None: + loss = self.loss_def.get_loss(output, model) # type: ignore self.value.update(loss) # type: ignore +class L2RegLogError(LossError): + """ + Logs (legacy) ||w||^2 of L2 regularization, i.e. 2 * get_loss + """ + + def update( + self, output: 'AtomGraphData', model: Optional[Callable] = None + ) -> None: + loss = 2.0 * self.loss_def.get_loss(output, model) # type: ignore + self.value.update(loss) # type: ignore + + +class ModalWeightCosine(ErrorMetric): + """ + Cosine similarity between modal weight views. + This is an indicator, not a loss. + """ + + def __init__( + self, + name: str, + loss_def: LossDefinition, + **kwargs, + ) -> None: + super().__init__( + name, + ignore_unlabeled=loss_def.ignore_unlabeled, + **kwargs, + ) + self.loss_def = loss_def + + def update( + self, output: 'AtomGraphData', model: Optional[Callable] = None + ) -> None: + cosine = self.loss_def.get_cosine(output, model) # type: ignore + self.value.update(cosine) # type: ignore + + class CombinedError(ErrorMetric): """ Combine multiple error metrics with weights @@ -288,9 +357,11 @@ def __init__(self, metrics: List[Tuple[ErrorMetric, float]], **kwargs) -> None: self.metrics = metrics assert kwargs['unit'] is None - def update(self, output: 'AtomGraphData') -> None: + def update( + self, output: 'AtomGraphData', model: Optional[Callable] = None + ) -> None: for metric, _ in self.metrics: - metric.update(output) + metric.update(output, model) def reset(self) -> None: for metric, _ in self.metrics: @@ -323,16 +394,16 @@ def __init__(self, metrics: List[ErrorMetric]) -> None: self.history = [] self.metrics = metrics - def _update(self, output: 'AtomGraphData') -> None: + def _update(self, output: 'AtomGraphData', model=None) -> None: for metric in self.metrics: - metric.update(output) + metric.update(output, model) - def update(self, output: 'AtomGraphData', no_grad=True) -> None: + def update(self, output: 'AtomGraphData', no_grad=True, model=None) -> None: if no_grad: with torch.no_grad(): - self._update(output) + self._update(output, model) else: - self._update(output) + self._update(output, model) def get_metric_dict(self, with_unit=True) -> Dict[str, float]: return {metric.key_str(with_unit): metric.get() for metric in self.metrics} @@ -409,9 +480,18 @@ def from_config( config: Dict[str, Any], loss_functions: Optional[List[Tuple[LossDefinition, float]]] = None, ) -> 'ErrorRecorder': - loss_cls = loss_dict[config.get(KEY.LOSS, 'mse').lower()] - loss_param = config.get(KEY.LOSS_PARAM, {}) - criteria = loss_cls(**loss_param) if loss_functions is None else None + loss_config = config.get(KEY.LOSS, 'mse') + if isinstance(loss_config, dict) and loss_functions is None: + raise NotImplementedError( + 'Structured loss config is not supported in train_v1. ' + 'Use train_v2 instead.' + ) + if loss_functions is None: + loss_cls = loss_dict[loss_config.lower()] + loss_param = config.get(KEY.LOSS_PARAM, {}) + criteria = loss_cls(**loss_param) + else: + criteria = None err_config = config.get(KEY.ERROR_RECORD, False) if not err_config: @@ -436,6 +516,22 @@ def from_config( ) ) continue + elif err_type == 'Modal_cos': # special case + metric_cls = ModalWeightCosine + metric_kwargs['loss_def'], _ = _get_loss_function_from_name( + loss_functions, 'L2_modal' + ) + metric_kwargs.pop('unit', None) + err_metrics.append(metric_cls(**metric_kwargs)) + continue + elif err_type == 'L2_modal': # special case + metric_kwargs['loss_def'], _ = _get_loss_function_from_name( + loss_functions, 'L2_modal' + ) + metric_kwargs.pop('unit', None) + metric_kwargs['name'] += f'_{metric_name}' + err_metrics.append(L2RegLogError(**metric_kwargs)) + continue metric_cls = ErrorRecorder.METRIC_DICT[metric_name] assert isinstance(metric_kwargs['name'], str) if metric_name == 'Loss': diff --git a/sevenn/model_build.py b/sevenn/model_build.py index 6975e2c1..f0482f11 100644 --- a/sevenn/model_build.py +++ b/sevenn/model_build.py @@ -144,7 +144,15 @@ def init_shift_scale( ) -> Union[Rescale, SpeciesWiseRescale, ModalWiseRescale]: # for mm, ex, shift: modal_idx -> shifts shift_scale = [] - train_shift_scale = config[KEY.TRAIN_SHIFT_SCALE] + train_shift = config.get(KEY.TRAIN_SHIFT, False) + train_scale = config.get(KEY.TRAIN_SCALE, False) + + # Legacy: train_shift_scale overrides both + # TODO: log this as legacy warning + train_shift_scale = config.get(KEY.TRAIN_SHIFT_SCALE, False) + if train_shift_scale: + train_shift = True + train_scale = True type_map = config[KEY.TYPE_MAP] # in case of modal, shift or scale has more dims [][] @@ -159,6 +167,7 @@ def init_shift_scale( shift_scale.append(s) shift, scale = shift_scale + ss_kwargs = {'train_shift': train_shift, 'train_scale': train_scale} rescale_module = None if config.get(KEY.USE_MODALITY, False): rescale_module = ModalWiseRescale.from_mappers( # type: ignore @@ -168,13 +177,13 @@ def init_shift_scale( config[KEY.USE_MODAL_WISE_SCALE], type_map=type_map, modal_map=config[KEY.MODAL_MAP], - train_shift_scale=train_shift_scale, + **ss_kwargs, ) elif all([isinstance(s, float) for s in shift_scale]): - rescale_module = Rescale(shift, scale, train_shift_scale=train_shift_scale) + rescale_module = Rescale(shift, scale, **ss_kwargs) elif any([isinstance(s, list) for s in shift_scale]): rescale_module = SpeciesWiseRescale.from_mappers( # type: ignore - shift, scale, type_map=type_map, train_shift_scale=train_shift_scale + shift, scale, type_map=type_map, **ss_kwargs ) else: raise ValueError('shift, scale should be list of float or float') diff --git a/sevenn/nn/scale.py b/sevenn/nn/scale.py index c8cbe87b..551355d4 100644 --- a/sevenn/nn/scale.py +++ b/sevenn/nn/scale.py @@ -30,16 +30,21 @@ def __init__( scale: float, data_key_in: str = KEY.SCALED_ATOMIC_ENERGY, data_key_out: str = KEY.ATOMIC_ENERGY, + train_shift: bool = False, + train_scale: bool = False, train_shift_scale: bool = False, **kwargs, ) -> None: assert isinstance(shift, float) and isinstance(scale, float) super().__init__() + if train_shift_scale: + train_shift = True + train_scale = True self.shift = nn.Parameter( - torch.FloatTensor([shift]), requires_grad=train_shift_scale + torch.FloatTensor([shift]), requires_grad=train_shift ) self.scale = nn.Parameter( - torch.FloatTensor([scale]), requires_grad=train_shift_scale + torch.FloatTensor([scale]), requires_grad=train_scale ) self.key_input = data_key_in self.key_output = data_key_out @@ -71,9 +76,14 @@ def __init__( data_key_in: str = KEY.SCALED_ATOMIC_ENERGY, data_key_out: str = KEY.ATOMIC_ENERGY, data_key_indices: str = KEY.ATOM_TYPE, + train_shift: bool = False, + train_scale: bool = False, train_shift_scale: bool = False, ) -> None: super().__init__() + if train_shift_scale: + train_shift = True + train_scale = True assert isinstance(shift, float) or isinstance(shift, list) assert isinstance(scale, float) or isinstance(scale, list) @@ -95,10 +105,10 @@ def __init__( scale = [scale] * num_species if isinstance(scale, float) else scale self.shift = nn.Parameter( - torch.FloatTensor(shift), requires_grad=train_shift_scale + torch.FloatTensor(shift), requires_grad=train_shift ) self.scale = nn.Parameter( - torch.FloatTensor(scale), requires_grad=train_shift_scale + torch.FloatTensor(scale), requires_grad=train_scale ) self.key_input = data_key_in self.key_output = data_key_out @@ -180,14 +190,19 @@ def __init__( data_key_atom_indices: str = KEY.ATOM_TYPE, use_modal_wise_shift: bool = False, use_modal_wise_scale: bool = False, + train_shift: bool = False, + train_scale: bool = False, train_shift_scale: bool = False, ) -> None: super().__init__() + if train_shift_scale: + train_shift = True + train_scale = True self.shift = nn.Parameter( - torch.FloatTensor(shift), requires_grad=train_shift_scale + torch.FloatTensor(shift), requires_grad=train_shift ) self.scale = nn.Parameter( - torch.FloatTensor(scale), requires_grad=train_shift_scale + torch.FloatTensor(scale), requires_grad=train_scale ) self.key_input = data_key_in self.key_output = data_key_out diff --git a/sevenn/parse_input.py b/sevenn/parse_input.py index 099e4c03..d302431b 100644 --- a/sevenn/parse_input.py +++ b/sevenn/parse_input.py @@ -186,6 +186,11 @@ def init_data_config(config: Dict[str, Any]) -> Dict[str, Any]: else: data_meta[load_data_key] = False + for k in config: + if k.startswith('load_') and k.endswith('_sequence'): + # TODO: check validity in here, not `from_config` + data_meta[k] = config[k] + for key, default in _const.DEFAULT_DATA_CONFIG.items(): data_meta[key] = config_initialize( key, config, default, _const.DATA_CONFIG_CONDITION diff --git a/sevenn/scripts/convert_model_modality.py b/sevenn/scripts/convert_model_modality.py index deca42a1..f0025775 100644 --- a/sevenn/scripts/convert_model_modality.py +++ b/sevenn/scripts/convert_model_modality.py @@ -5,15 +5,11 @@ import torch.nn as nn from e3nn.o3 import Irreps, Linear +import sevenn._const as CONST import sevenn._keys as KEY from sevenn.model_build import build_E3_equivariant_model -modal_module_dict = { - KEY.USE_MODAL_NODE_EMBEDDING: 'onehot_to_feature_x', - KEY.USE_MODAL_SELF_INTER_INTRO: 'self_interaction_1', - KEY.USE_MODAL_SELF_INTER_OUTRO: 'self_interaction_2', - KEY.USE_MODAL_OUTPUT_BLOCK: 'reduce_input_to_hidden', -} +modal_module_dict = CONST.IMPLEMENTED_MODAL_MODULE_DICT def _get_scalar_index(irreps: Irreps) -> List[int]: diff --git a/sevenn/scripts/processing_by_batch.py b/sevenn/scripts/processing_by_batch.py new file mode 100644 index 00000000..92534d7d --- /dev/null +++ b/sevenn/scripts/processing_by_batch.py @@ -0,0 +1,258 @@ +import math +import os +import time +from copy import deepcopy +from typing import Any, Dict, Optional + +import numpy as np + +import sevenn._keys as KEY +from sevenn.error_recorder import AverageNumber, ErrorRecorder +from sevenn.logger import Logger +from sevenn.train.trainer import Trainer +from sevenn.util import unique_filepath + + +def update_config_for_batch_training(config: Dict[str, Any], loaders) -> None: + """ + Called from sevenn/scripts/train.py, between dataset build and model build + Update scheduler parameters for batch-level training. + + This converts epoch-based scheduler parameters to step-based parameters + when using batch training mode. + """ + train_loader = loaders['trainset'] + # convert float type `epoch` related parameters for batch training + effective_batch_size = config[KEY.WORLD_SIZE] * config[KEY.BATCH_SIZE] + steps_per_epoch = math.ceil( + train_loader.sampler.total_size / effective_batch_size + ) + + scheduler_type = config.get(KEY.SCHEDULER, 'exponentiallr').lower() + scheduler_param = config.get(KEY.SCHEDULER_PARAM, {}) + config[KEY.SCHEDULER_BATCH_MODE] = scheduler_param.pop( + KEY.SCHEDULER_BATCH_MODE, False + ) + + if scheduler_type == 'onecyclelr': # special case, always batch mode + total_steps = scheduler_param.get('total_steps', None) + if total_steps is None: + # total_steps not given, automatically calculated + # allow epochs to be float for SWA + epochs = scheduler_param.get('epochs', None) + if epochs is None: + raise ValueError('One of total_steps or epochs should be given') + total_steps = math.ceil(epochs * steps_per_epoch) + config[KEY.SCHEDULER_PARAM]['total_steps'] = total_steps + config[KEY.SCHEDULER_BATCH_MODE] = True + + elif config[KEY.SCHEDULER_BATCH_MODE]: + scheduler_epoch_params = { + 'linearlr': ['total_iters', lambda x, y: math.ceil(x * y)], + 'cosineannealinglr': ['T_max', lambda x, y: math.ceil(x * y)], + 'exponentiallr': ['gamma', lambda x, y: x ** (1 / y)], + }.get(scheduler_type, None) + if scheduler_epoch_params is None: + raise NotImplementedError( + f'Scheduler batch mode not implemented for {scheduler_type}.' + ) + + config[KEY.SCHEDULER_PARAM][scheduler_epoch_params[0]] = ( + scheduler_epoch_params[1]( + config[KEY.SCHEDULER_PARAM][scheduler_epoch_params[0]], + steps_per_epoch, + ) + ) + + +def processing_by_batch( + config: dict, + trainer: Trainer, + loaders: dict, + data_progress: dict, + start_epoch: int = 1, + train_loader_key: str = 'trainset', + error_recorder: Optional[ErrorRecorder] = None, + total_epoch: Optional[int] = None, + per_epoch: Optional[float] = None, + best_metric_loader_key: str = 'validset', + best_metric: Optional[str] = None, + write_csv: bool = True, + working_dir: Optional[str] = None, +): + """ + Batch-level training loop for large-scale training. + + Unlike epoch-level training (processing_epoch_v2), this function: + - Saves checkpoints at configurable intervals within an epoch + - Supports resuming from exact data position + - Uses OrderedSampler for deterministic DDP training + + Args: + config: Training configuration dictionary + trainer: Trainer instance + loaders: Dict of dataloaders (must include 'trainset') + data_progress: Dict tracking data iteration progress + start_epoch: Starting epoch number + train_loader_key: Key for training dataloader + error_recorder: ErrorRecorder instance (created from config if None) + total_epoch: Total epochs to train (from config if None) + per_epoch: Checkpoint frequency as fraction of epoch (from config if None) + best_metric_loader_key: Key for validation loader used for best metric + best_metric: Metric name for tracking best model + write_csv: Whether to write CSV log + working_dir: Working directory for checkpoints + """ + log = Logger() + write_csv = write_csv and log.rank == 0 + working_dir = working_dir or os.getcwd() + prefix = f'{os.path.abspath(working_dir)}/' + + total_epoch = total_epoch or config[KEY.EPOCH] + per_epoch = per_epoch or config.get(KEY.PER_EPOCH, 0.1) + best_metric = best_metric or config.get(KEY.BEST_METRIC, 'TotalLoss') + recorder = error_recorder or ErrorRecorder.from_config( + config, trainer.loss_functions + ) + recorders = {k: deepcopy(recorder) for k in loaders} + + best_val = float('inf') + best_key = None + if best_metric_loader_key in recorders: + best_key = recorders[best_metric_loader_key].get_key_str(best_metric) + if best_key is None: + log.writeline( + f'Failed to get error recorder key: {best_metric} or ' + + f'{best_metric_loader_key} is missing. There will be no best ' + + 'checkpoint.' + ) + + csv_path = unique_filepath(f'{prefix}/lc.csv') + if write_csv: + head = ['epoch', 'batch', 'lr'] + for k, rec in recorders.items(): + head.extend(list(rec.get_dct(prefix=k))) + with open(csv_path, 'w') as f: + f.write(','.join(head) + '\n') + + train_loader = loaders['trainset'] + + if data_progress[KEY.TOTAL_DATA_NUM] < 0: # fresh start or reset + data_progress[KEY.TOTAL_DATA_NUM] = len(train_loader.sampler.sequence) + else: # continue + train_loader.sampler.continue_from_data_progress(**data_progress) + + effective_batch_size = config[KEY.WORLD_SIZE] * config[KEY.BATCH_SIZE] + start_batch = (data_progress[KEY.CURRENT_DATA_IDX]) // effective_batch_size + total_step = data_progress[KEY.TOTAL_DATA_NUM] // effective_batch_size + if data_progress[KEY.TOTAL_DATA_NUM] % effective_batch_size != 0: + total_step += 1 + save_per_epoch = int(1 / per_epoch) + save_batch_idx = np.linspace(1, total_step, save_per_epoch + 1) + save_batch_idx = [int(idx) for idx in save_batch_idx[1:]] + + if data_progress[KEY.TOTAL_DATA_NUM] == data_progress[KEY.CURRENT_DATA_IDX]: + start_epoch += 1 # continuing from end of epoch + start_batch = 0 + train_loader.sampler.permutate_sequence() # update rng state + train_loader.sampler.refresh_sequence() # make starting index to zero + + trainer.write_checkpoint( + f'{prefix}/checkpoint_initial.pth', + config=config, + epoch=0, + data_progress=data_progress, + ) + + scheduler_update_every_batch = ( + config.get(KEY.SCHEDULER_BATCH_MODE, False) + ) + + # TODO: too long, refactor more + log.writeline('Entering training loop') + for epoch in range(start_epoch, total_epoch + 1): # one indexing + data_progress[KEY.NUMPY_RNG_STATE] = train_loader.sampler.get_rng_state() + log.timer_start('epoch') + log.timer_start('batch') + + dl_timing = AverageNumber() + dl_end = time.time() + for idx, batch in enumerate(train_loader): + dl_timing.update(time.time() - dl_end) + current_batch_idx = idx + 1 + if epoch == start_epoch: + current_batch_idx += start_batch # continuing from middle of epoch + save = current_batch_idx in save_batch_idx + + if save: + lr = trainer.get_lr() + log.bar() + log.writeline( + f'Epoch {epoch}/{total_epoch} ' + + f'Batch {current_batch_idx}/{total_step} lr: {lr:8f}' + ) + log.bar() + + trainer.train_one_batch(batch, recorders[train_loader_key]) + if scheduler_update_every_batch: # onecyclelr + trainer.scheduler_step(best_val) + + if save: + csv_dct = { + 'epoch': str(epoch), + 'batch': str(current_batch_idx), + 'lr': f'{trainer.get_lr():8f}', + } + errors = {} + for k, loader in loaders.items(): + rec = recorders[k] + if k != train_loader_key: + print('valid', k, flush=True) + trainer.run_one_epoch(loader, False, rec) + if trainer.distributed: + trainer.recorder_all_reduce(rec) + csv_dct.update(rec.get_dct(prefix=k)) + errors[k] = rec.epoch_forward() + log.write_full_table(list(errors.values()), list(errors)) + + batch_name = ( + f'_{save_batch_idx.index(current_batch_idx) + 1}' + if current_batch_idx != save_batch_idx[-1] + else '' + ) + data_progress[KEY.CURRENT_DATA_IDX] = min( + current_batch_idx * effective_batch_size, + data_progress[KEY.TOTAL_DATA_NUM], + ) + trainer.write_checkpoint( + f'{prefix}/checkpoint_{epoch}{batch_name}.pth', + config=config, + epoch=epoch, + data_progress=data_progress, + ) + + if write_csv: + with open(csv_path, 'a') as f: + f.write(','.join(list(csv_dct.values())) + '\n') + + if best_key and errors[best_metric_loader_key][best_key] < best_val: + trainer.write_checkpoint( + f'{prefix}/checkpoint_best.pth', config=config, epoch=epoch + ) + best_val = errors[best_metric_loader_key][best_key] + log.writeline('Best checkpoint written') + + log.timer_end('batch', message=f'Batch {current_batch_idx} elapsed') + if config[KEY.IS_DDP]: + dl_timing._ddp_reduce(trainer.device) + log.writeline(f'data loading, per (sec): {dl_timing.get():.4f}') + log.writeline(f'data loading, sum *ALL* (sec): {dl_timing._sum:.4f}') + dl_timing = AverageNumber() + log.timer_start('batch') + dl_end = time.time() + # batch loop indent + + if not scheduler_update_every_batch: + trainer.scheduler_step(best_val) + log.timer_end('epoch', message=f'Epoch {epoch} elapsed') + return trainer diff --git a/sevenn/scripts/processing_continue.py b/sevenn/scripts/processing_continue.py index 169bb362..6bef8073 100644 --- a/sevenn/scripts/processing_continue.py +++ b/sevenn/scripts/processing_continue.py @@ -1,6 +1,6 @@ import os import warnings -from typing import Any, Dict, List, Tuple +from typing import Any, Dict import torch @@ -13,15 +13,18 @@ ) -def processing_continue_v2( - config: Dict[str, Any], -) -> Tuple[List[Dict[str, torch.Tensor]], int]: # simpler +# TODO: check backward compatibility +def processing_continue_v2(config: Dict[str, Any]): """ Replacement of processing_continue, Skips model compatibility + + Returns: + (state_dicts, epoch, data_progress or None) """ log = Logger() continue_dct = config[KEY.CONTINUE] + train_by_batch = config.get(KEY.TRAIN_BY_BATCH, False) log.write('\nContinue found, loading checkpoint\n') checkpoint = util.load_checkpoint(continue_dct[KEY.CHECKPOINT]) @@ -79,17 +82,37 @@ def processing_continue_v2( from_epoch = checkpoint.epoch or 0 log.writeline(f'Checkpoint previous epoch was: {from_epoch}') - epoch = 1 if continue_dct[KEY.RESET_EPOCH] else from_epoch + 1 - log.writeline(f'epoch start from {epoch}') - log.writeline('checkpoint loading successful') + # For batch training, don't increment epoch yet (handled by processing_by_batch) + if train_by_batch: + epoch = 1 if continue_dct[KEY.RESET_EPOCH] else from_epoch + else: + epoch = 1 if continue_dct[KEY.RESET_EPOCH] else from_epoch + 1 + log.writeline(f'epoch start from {epoch}') state_dicts = [ model_state_dict_cp, optimizer_state_dict_cp, scheduler_state_dict_cp, ] - return state_dicts, epoch + + # Handle data progress for batch training + data_progress = {} + if train_by_batch: + data_progress = { + KEY.TOTAL_DATA_NUM: -1, + KEY.CURRENT_DATA_IDX: 0, + KEY.NUMPY_RNG_STATE: None, + } + if hasattr(checkpoint, 'data_progress') and checkpoint.data_progress: + if not continue_dct[KEY.RESET_DATA_PROGRESS]: + data_progress.update(checkpoint.data_progress) + log.writeline(f'epoch start from {epoch}') + log.writeline(f'data index start from {data_progress[KEY.CURRENT_DATA_IDX]}') # noqa: E501 + + log.writeline('checkpoint loading success') + + return state_dicts, epoch, data_progress or {} def check_config_compatible(config: Dict[str, Any], config_cp: Dict[str, Any]): @@ -125,7 +148,12 @@ def check_config_compatible(config: Dict[str, Any], config_cp: Dict[str, Any]): except KeyError: return - TRAINABLE_CONFIGS = [KEY.TRAIN_DENOMINTAOR, KEY.TRAIN_SHIFT_SCALE] + TRAINABLE_CONFIGS = [ + KEY.TRAIN_DENOMINTAOR, + KEY.TRAIN_SHIFT_SCALE, + KEY.TRAIN_SHIFT, + KEY.TRAIN_SCALE, + ] if ( any((not cntdct[KEY.RESET_SCHEDULER], not cntdct[KEY.RESET_OPTIMIZER])) and all(config[k] == config_cp[k] for k in TRAINABLE_CONFIGS) is False diff --git a/sevenn/scripts/train.py b/sevenn/scripts/train.py index 06eba2e7..8eb1dd5a 100644 --- a/sevenn/scripts/train.py +++ b/sevenn/scripts/train.py @@ -1,5 +1,8 @@ +import importlib.util +from pathlib import Path from typing import Any, Dict, List, Optional +import numpy as np import torch.distributed as dist from torch.utils.data.dataset import Dataset from torch.utils.data.distributed import DistributedSampler @@ -11,43 +14,111 @@ from sevenn.scripts.processing_continue import ( convert_modality_of_checkpoint_state_dct, ) -from sevenn.train.reewc import ( - ReewcTrainer, +from sevenn.train.reewc.rehearsal import ( build_memory_loader, reewc_dataset_keys, validate_reewc_config, ) +from sevenn.train.reewc.trainer import ReewcTrainer +from sevenn.train.sampler import OrderedSampler from sevenn.train.trainer import Trainer def loader_from_config( - config: Dict[str, Any], dataset: Dataset, is_train: bool = False + config: Dict[str, Any], + dataset: Dataset, + dataset_key: str, ) -> DataLoader: + """ + Create DataLoader from config. + + Args: + config: Configuration dictionary + dataset: Dataset to create loader for + (or dict with 'dataset' and 'batch_size') + dataset_key: Key identifying the dataset + """ batch_size = config[KEY.BATCH_SIZE] - shuffle = is_train and config[KEY.TRAIN_SHUFFLE] + + if isinstance(dataset, dict): + batch_size = dataset.get('batch_size', batch_size) + dataset = dataset['dataset'] + + shuffle = config[KEY.TRAIN_SHUFFLE] sampler = None + loader_args = {'dataset': dataset, 'batch_size': batch_size, 'shuffle': shuffle} if KEY.NUM_WORKERS in config and config[KEY.NUM_WORKERS] > 0: loader_args.update({'num_workers': config[KEY.NUM_WORKERS]}) + if (loader_kwargs := config.get(KEY.LOADER_KWARGS, None)) is not None: + loader_args.update(**loader_kwargs) + + world_size, rank = 1, 0 if config[KEY.IS_DDP]: dist.barrier() - sampler = DistributedSampler( - dataset, dist.get_world_size(), dist.get_rank(), shuffle=shuffle - ) + world_size = dist.get_world_size() + rank = dist.get_rank() + sampler = DistributedSampler(dataset, world_size, rank, shuffle=shuffle) loader_args.update({'sampler': sampler}) loader_args.pop('shuffle') # sampler is mutually exclusive with shuffle + + # Use OrderedSampler for batch training mode to preserve data order + # verified only for validset + # TODO: I think 'train_by_batch' and 'sampling validset' is independent, + # so 'train_by_batch' should be removed + if config.get(KEY.TRAIN_BY_BATCH, False): + sequence = config.get(f'load_{dataset_key}_sequence', {}).get( + 'total_sequence_path', None + ) + + sampler = OrderedSampler( + dataset=dataset, + sequence=np.load(sequence) if sequence else None, + shuffle=shuffle, + seed=config.get(KEY.RANDOM_SEED, 777), + world_size=world_size, + rank=rank, + ) + # sampler is mutually exclusive with shuffle + loader_args.update({'sampler': sampler, 'shuffle': None}) return DataLoader(**loader_args) +def datasets_from_py(config, script): + if isinstance(script, list): + assert len(script) == 1, 'Need single python script' + script = script[0] + + file_path = Path(script).resolve() + print(f'Init dataset from {file_path}', flush=True) + spec = importlib.util.spec_from_file_location('dataset', file_path) + module = importlib.util.module_from_spec(spec) # type: ignore + spec.loader.exec_module(module) # type: ignore + + ret = module.dataset(config) + assert isinstance(ret, dict) and 'trainset' in ret + return ret + + +# TODO: check backward compatibility this part (batch vs. epoch) def train_v2(config: Dict[str, Any], working_dir: str) -> None: """ Main program flow, since v0.9.6 + + Supports: + - Epoch-level training (default) + - Batch-level training (train_by_batch: true) """ + import sevenn.train.aselmdb_dataset as aselmdb_dataset import sevenn.train.atoms_dataset as atoms_dataset import sevenn.train.graph_dataset as graph_dataset import sevenn.train.modal_dataset as modal_dataset + from .processing_by_batch import ( + processing_by_batch, + update_config_for_batch_training, + ) from .processing_continue import processing_continue_v2 from .processing_epoch import processing_epoch_v2 @@ -62,47 +133,80 @@ def train_v2(config: Dict[str, Any], working_dir: str) -> None: config[KEY.LOAD_TRAINSET] = config.pop(KEY.LOAD_DATASET) validate_reewc_config(config) + # Initialize data progress for batch training + train_by_batch = config.get(KEY.TRAIN_BY_BATCH, False) + + data_progress = {} + if train_by_batch: + data_progress = { + KEY.TOTAL_DATA_NUM: -1, + KEY.CURRENT_DATA_IDX: 0, + KEY.NUMPY_RNG_STATE: None, + } # config updated start_epoch = 1 state_dicts: Optional[List[dict]] = None if config[KEY.CONTINUE][KEY.CHECKPOINT]: - state_dicts, start_epoch = processing_continue_v2(config) + # data_progress is non-empty only if train_by_batch is True + state_dicts, start_epoch, data_progress = processing_continue_v2(config) - if config.get(KEY.USE_MODALITY, False): + # Load datasets based on type + dataset_type = config[KEY.DATASET_TYPE] + if ( + config.get(KEY.USE_MODALITY, False) + and not config[KEY.DATASET_TYPE] == 'custom' + ): datasets = modal_dataset.from_config(config, working_dir) - elif config[KEY.DATASET_TYPE] == 'graph': + elif dataset_type == 'graph': datasets = graph_dataset.from_config( config, working_dir, dataset_keys=reewc_dataset_keys(config) ) - elif config[KEY.DATASET_TYPE] == 'atoms': + elif dataset_type == 'atoms': datasets = atoms_dataset.from_config(config, working_dir) + elif dataset_type == 'aselmdb': + datasets = aselmdb_dataset.from_config(config, working_dir) + elif dataset_type == 'custom': + datasets = datasets_from_py(config, config.get('load_trainset_path')) else: - raise ValueError(f'Unknown dataset type: {config[KEY.DATASET_TYPE]}') + raise ValueError(f'Unknown dataset type: {dataset_type}') + loaders = { - k: loader_from_config(config, v, is_train=(k == 'trainset')) - for k, v in datasets.items() + k: loader_from_config(config, v, dataset_key=k) for k, v in datasets.items() } - rehearsal = config.get(KEY.REHEARSAL, False) - memory_loader = build_memory_loader(config) if rehearsal else None + # Update scheduler config for batch training + if train_by_batch: + update_config_for_batch_training(config, loaders) log.write('\nModel building...\n') model = build_E3_equivariant_model(config) log.print_model_info(model, config) - if memory_loader is not None: + if config.get(KEY.REHEARSAL, False): + memory_loader = build_memory_loader(config) trainer = ReewcTrainer.from_config( model, config, memory_loader=memory_loader ) else: trainer = Trainer.from_config(model, config) + if state_dicts: trainer.load_state_dicts(*state_dicts, strict=False) - processing_epoch_v2( - config, trainer, loaders, start_epoch, working_dir=working_dir - ) + if train_by_batch: + processing_by_batch( + config, + trainer, + loaders, + data_progress, + start_epoch, + working_dir=working_dir, + ) + else: + processing_epoch_v2( + config, trainer, loaders, start_epoch, working_dir=working_dir + ) log.timer_end('total', message='Total wall time') @@ -128,8 +232,7 @@ def train(config, working_dir: str): train, valid, _ = processing_dataset(config, working_dir) datasets = {'dataset': train, 'validset': valid} loaders = { - k: loader_from_config(config, v, is_train=(k == 'dataset')) - for k, v in datasets.items() + k: loader_from_config(config, v, dataset_key=k) for k, v in datasets.items() } loaders = list(loaders.values()) diff --git a/sevenn/train/aselmdb_dataset.py b/sevenn/train/aselmdb_dataset.py new file mode 100644 index 00000000..b425e324 --- /dev/null +++ b/sevenn/train/aselmdb_dataset.py @@ -0,0 +1,740 @@ +from __future__ import annotations + +import bisect +import os +import os.path as osp +import time +import warnings +import zlib +from collections import Counter +from glob import glob +from pathlib import Path +from typing import Any, Callable, Dict, List, Optional, Union + +import ase +import lmdb +import numpy as np +import orjson +from ase.data import chemical_symbols +from ase.db.core import Database, now, ops +from ase.db.row import AtomsRow +from tqdm import tqdm + +import sevenn._keys as KEY +import sevenn.util as util +from sevenn._const import NUM_UNIV_ELEMENT +from sevenn.atom_graph_data import AtomGraphData +from sevenn.train.atoms_dataset import SevenNetAtomsDataset +from sevenn.train.dataload import _set_atoms_y + + +class LMDBDatabase(Database): + """ + Same class develop by Meta + + Copyright (c) Meta, Inc. and its affiliates. + + This source code is modified from the ASE db json backend + and is thus licensed under the corresponding LGPL2.1 license + + The ASE notice for the LGPL2.1 license is available here: + https://gitlab.com/ase/ase/-/blob/master/LICENSE + """ + + def __init__( + self, + filename: str | Path | None = None, + create_indices: bool = True, + use_lock_file: bool = False, + serial: bool = False, + readonly: bool = False, + *args, + **kwargs, + ) -> None: + """ + For the most part, this is identical to the standard ase db initiation + arguments, except that we add a readonly flag. + """ + super().__init__( + Path(filename), + create_indices, + use_lock_file, + serial, + *args, + **kwargs, + ) + + # Add a readonly mode for when we're only training + # to make sure there's no parallel locks + self.readonly = readonly + + if self.readonly: + # Open a new env + self.env = lmdb.open( + str(self.filename), + subdir=False, + meminit=False, + map_async=True, + readonly=True, + lock=False, + ) + + # Open a transaction and keep it open for fast read/writes! + self.txn = self.env.begin(write=False) + + else: + # Open a new env with write access + self.env = lmdb.open( + str(self.filename), + map_size=1099511627776 * 2, + subdir=False, + meminit=False, + map_async=True, + ) + + self.txn = self.env.begin(write=True) + + # Load all ids based on keys in the DB. + self.ids = [] + self.deleted_ids = [] + self._load_ids() + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, tb) -> None: + self.close() + + def close(self) -> None: + # Close the lmdb environment and transaction + self.txn.commit() + self.env.close() + + def _write( + self, + atoms: ase.Atoms | AtomsRow, + key_value_pairs: dict, + data: dict | None, + idx: int | None = None, + ) -> None: + Database._write(self, atoms, key_value_pairs, data) + + mtime = now() + + if isinstance(atoms, AtomsRow): + row = atoms + else: + row = AtomsRow(atoms) + row.ctime = mtime + row.user = os.getenv('USER') + + dct = {} + for key in row.__dict__: + if key[0] == '_' or key in row._keys or key == 'id': + continue + dct[key] = row[key] + + dct['mtime'] = mtime + + if key_value_pairs: + dct['key_value_pairs'] = key_value_pairs + + if data: + dct['data'] = data + + constraints = row.get('constraints') + if constraints: + dct['constraints'] = [constraint.todict() for constraint in constraints] + + # json doesn't like Cell objects, so make it an array + dct['cell'] = np.asarray(dct['cell']) + + if idx is None: + idx = self._nextid + nextid = idx + 1 + else: + data = self.txn.get(f"{idx}".encode('ascii')) + assert data is not None + + # Add the new entry + self.txn.put( + f"{idx}".encode('ascii'), + zlib.compress(orjson.dumps(dct, option=orjson.OPT_SERIALIZE_NUMPY)), + ) + # only append if idx is not in ids + if idx not in self.ids: + self.ids.append(idx) + self.txn.put( + 'nextid'.encode('ascii'), + zlib.compress( + orjson.dumps(nextid, option=orjson.OPT_SERIALIZE_NUMPY) + ), + ) + # check if id is in removed ids and remove accordingly + if idx in self.deleted_ids: + self.deleted_ids.remove(idx) + self._write_deleted_ids() + + return idx + + def _update( + self, + idx: int, + key_value_pairs: dict | None = None, + data: dict | None = None, + ): + # hack this to play nicely with ASE code + row = self._get_row(idx, include_data=True) + if data is not None or key_value_pairs is not None: + self._write( + atoms=row, idx=idx, key_value_pairs=key_value_pairs, data=data + ) + + def _write_deleted_ids(self): + self.txn.put( + 'deleted_ids'.encode('ascii'), + zlib.compress( + orjson.dumps(self.deleted_ids, option=orjson.OPT_SERIALIZE_NUMPY) + ), + ) + + def delete(self, ids: list[int]) -> None: + for idx in ids: + self.txn.delete(f"{idx}".encode('ascii')) + self.ids.remove(idx) + + self.deleted_ids += ids + self._write_deleted_ids() + + def _get_row(self, idx: int, include_data: bool = True): + if idx is None: + assert len(self.ids) == 1 + idx = self.ids[0] + data = self.txn.get(f"{idx}".encode('ascii')) + + if data is not None: + dct = orjson.loads(zlib.decompress(data)) + else: + raise KeyError(f"Id {idx} missing from the database!") + + if not include_data: + dct.pop('data', None) + + dct['id'] = idx + return AtomsRow(dct) + + def _get_row_by_index(self, index: int, include_data: bool = True): + """Auxiliary function to get the ith entry, rather than a specific id""" + data = self.txn.get(f"{self.ids[index]}".encode('ascii')) + + if data is not None: + dct = orjson.loads(zlib.decompress(data)) + else: + raise KeyError(f"Id {id} missing from the database!") + + if not include_data: + dct.pop('data', None) + + dct['id'] = id + return AtomsRow(dct) + + def _select( + self, + keys, + cmps: list[tuple[str, str, str]], + explain: bool = False, + verbosity: int = 0, + limit: int | None = None, + offset: int = 0, + sort: str | None = None, + include_data: bool = True, + columns: str = 'all', + ): + if explain: + yield {'explain': (0, 0, 0, 'scan table')} + return + + if sort is not None: + if sort[0] == '-': + reverse = True + sort = sort[1:] + else: + reverse = False + + rows = [] + missing = [] + for row in self._select(keys, cmps): + key = row.get(sort) + if key is None: + missing.append((0, row)) + else: + rows.append((key, row)) + + rows.sort(reverse=reverse, key=lambda x: x[0]) + rows += missing + + if limit: + rows = rows[offset : offset + limit] + for _, row in rows: + yield row + return + + if not limit: + limit = -offset - 1 + + cmps = [(key, ops[op], val) for key, op, val in cmps] + n = 0 + for idx in self.ids: + if n - offset == limit: + return + row = self._get_row(idx, include_data=include_data) + + for key in keys: + if key not in row: + break + else: + for key, op, val in cmps: + if isinstance(key, int): + value = np.equal(row.numbers, key).sum() + else: + value = row.get(key) + if key == 'pbc': + assert op in [ops['='], ops['!=']] + value = ''.join('FT'[x] for x in value) + if value is None or not op(value, val): + break + else: + if n >= offset: + yield row + n += 1 + + @property + def metadata(self): + """Load the metadata from the DB if present""" + if self._metadata is None: + metadata = self.txn.get('metadata'.encode('ascii')) + if metadata is None: + self._metadata = {} + else: + self._metadata = orjson.loads(zlib.decompress(metadata)) + + return self._metadata.copy() + + @metadata.setter + def metadata(self, dct): + self._metadata = dct + + # Put the updated metadata dictionary + self.txn.put( + 'metadata'.encode('ascii'), + zlib.compress(orjson.dumps(dct, option=orjson.OPT_SERIALIZE_NUMPY)), + ) + + @property + def _nextid(self): + """Get the id of the next row to be written""" + # Get the nextid + nextid_data = self.txn.get('nextid'.encode('ascii')) + return orjson.loads(zlib.decompress(nextid_data)) if nextid_data else 1 + + def count(self, selection=None, **kwargs) -> int: + """Count rows. + + See the select() method for the selection syntax. Use db.count() or + len(db) to count all rows. + """ + if selection is not None: + n = 0 + for _row in self.select(selection, **kwargs): + n += 1 + return n + else: + # Fast count if there's no queries! Just get number of ids + return len(self.ids) + + def _load_ids(self) -> None: + """Load ids from the DB + + Since ASE db ids are mostly 1-N integers, but can be missing entries + if ids have been deleted. To save space and operating under the assumption + that there will probably not be many deletions in most OCP datasets, + we just store the deleted ids. + """ + + # Load the deleted ids + deleted_ids_data = self.txn.get('deleted_ids'.encode('ascii')) + if deleted_ids_data is not None: + self.deleted_ids = orjson.loads(zlib.decompress(deleted_ids_data)) + + # Reconstruct the full id list + self.ids = [ + i for i in range(1, self._nextid) if i not in set(self.deleted_ids) + ] + + +class AseDBDataset: + """ + Modified, combined code for fairchem AseDBDataset w/o dependency + """ + def __init__( + self, + src: Union[List[str], str], + ): + if isinstance(src, list): + filepaths = [] + for path in src: + if os.path.isdir(path): + filepaths.extend(glob(f"{path}/*")) + elif os.path.isfile(path): + filepaths.append(path) + else: + raise RuntimeError(f"Error reading dataset in {path}!") + elif os.path.isfile(src): + filepaths = [src] + elif os.path.isdir(src): + filepaths = glob(f'{src}/*') + else: + filepaths = glob(src) + + self.dbs = [] + + for path in sorted(filepaths): + try: + self.dbs.append(self.connect_db(path)) + except ValueError: + pass + + # In order to get all of the unique IDs using the default ASE db interface + # we have to load all the data and check ids using a select. This is extremely # noqa: E501 + # inefficient for large dataset. If the db we're using already presents a list of # noqa: E501 + # ids and there is no query, we can just use that list instead and save ourselves # noqa: E501 + # a lot of time! + self.db_ids = [] + for db in self.dbs: + if hasattr(db, 'ids'): + self.db_ids.append(db.ids) + else: + # this is the slow alternative + self.db_ids.append([row.id for row in db.select()]) + + idlens = [len(ids) for ids in self.db_ids] + self._idlen_cumulative = np.cumsum(idlens).tolist() + + self.ids = list(range(sum(idlens))) + self.num_samples = len(self.ids) + + if self.num_samples == 0: + raise ValueError(f"No valid ase data found, check src {src}!") + + def __len__(self) -> int: + return self.num_samples + + def connect_db(self, path: str) -> Database: + if any(path.endswith(ext) for ext in ['aselmdb', 'lmdb']): + return LMDBDatabase(path, readonly=True) + + return ase.db.connect(path) + + def get_atoms(self, idx: int) -> ase.Atoms: + db_idx = bisect.bisect(self._idlen_cumulative, idx) + + el_idx = idx + if db_idx != 0: + el_idx = idx - self._idlen_cumulative[db_idx - 1] + assert el_idx >= 0 + + atoms_row = self.dbs[db_idx]._get_row(self.db_ids[db_idx][el_idx]) + atoms = atoms_row.toatoms() + + if isinstance(atoms_row.data, dict): + atoms.info.update(atoms_row.data) + + return atoms + + +class SevenNetASElmdbDataset(SevenNetAtomsDataset): + def __init__( + self, + cutoff: float, + files: Union[str, List[str]], + # sequence: Optional[List[int]] = None, + stat_sequence_info: Union[str, float, int] = 10000, + is_auto_mode: bool = False, + atoms_filter: Optional[Callable] = None, # not used yet + atoms_transform: Optional[Callable] = None, + graph_transform: Optional[Callable] = None, + **process_kwargs, + ): + self.cutoff = cutoff + if isinstance(files, str): + files = [files] # user convenience + files = [osp.abspath(file) for file in files] + + self._files = files + self.is_auto_mode = is_auto_mode + self.atoms_filter = atoms_filter + self.atoms_trasform = atoms_transform + self.graph_trasform = graph_transform + self._scanned = False + self._avg_num_neigh_approx = None + self.statistics = {} + + self._dataset = AseDBDataset(src=files) + + if isinstance(stat_sequence_info, str) and osp.exists(stat_sequence_info): + self.stat_sequence = np.load(stat_sequence_info) + + elif isinstance(stat_sequence_info, int): + sample_num = min(len(self), stat_sequence_info) + self.stat_sequence = np.random.choice( + np.arange(len(self)), sample_num, replace=False + ) + + elif isinstance(stat_sequence_info, float): + sample_num = int(len(self) * stat_sequence_info) + self.stat_sequence = np.random.choice( + np.arange(len(self)), sample_num, replace=False + ) + + else: + raise ValueError( + 'stat_sequence_info should be one of str, int, float, ' + + f'but got {type(stat_sequence_info)}' + ) + + def __len__(self): + # total, run_sequence deprecated. + # Should only be used in OrderedSampler.__init__ + # return len(self._run_sequence) + return len(self._dataset) + + def __getitem__(self, index): + # idx = self._run_sequence[index] + # atoms = self.set_atoms_y_with_idx(idx) + atoms = self.set_atoms_y_with_idx(index) + if self.atoms_trasform is not None: + atoms = self.atoms_trasform(atoms) + + graph = self._graph_build(atoms) + if self.graph_trasform is not None: + graph = self.graph_trasform(graph) + + return AtomGraphData.from_numpy_dict(graph) + + def preload(self): + print('Start preload..., fulfilling OS cache') + print('Quick if it was already loaded into cache. Check "free -h"') + print('May useless if the RAM is smaller than dataset size') + print(f'Number of files: {len(self._dataset.dbs)}', flush=True) + end = time.time() + for lmdbdb in tqdm(self._dataset.dbs): + lmdb_env = lmdbdb.env + with lmdb_env.begin(write=False) as txn: + with txn.cursor() as cursor: + if cursor.first(): # move cursor to start, skip 0 len db + while cursor.next(): + _, _ = cursor.key(), cursor.value() + # dct = orjson.loads(zlib.decompress(v)) + print(f'Preload elapsed (sec): {time.time() - end:.4f}', flush=True) + + def set_atoms_y_with_idx(self, idx): + atoms = self._dataset.get_atoms(idx) + atoms = _set_atoms_y([atoms])[0] + return atoms + + @property + def species(self): + mode = ( + 'total' if self.is_auto_mode else 'stat' + ) # species should be fully scanned only in `auto` mode + self.run_stat(mode=mode) + return [z for z in self.statistics['_natoms'].keys() if z != 'total'] + + @property + def avg_num_neigh(self, n_sample=10000): + if self._avg_num_neigh_approx is None: + if len(self.stat_sequence) > n_sample: + warnings.warn( + """SevenNetASElmdbDataset does not provide correct avg_num_neigh + as it does not build graph. We will compute only random 10000 + structures graph to approximate this value. If you want more + precise avg_num_neigh, use SevenNetGraphDataset. If it is not + viable due to memory limit, you need online algorithm to do this + , which is not yet implemented in the SevenNet""" + ) + n_sample = min(len(self.stat_sequence), n_sample) + indices = np.random.choice(self.stat_sequence, n_sample, replace=False) + n_neigh = [] + for i in indices: + atoms = self.set_atoms_y_with_idx(i) + graph = self._graph_build(atoms) + _, nn = np.unique(graph[KEY.EDGE_IDX][0], return_counts=True) + n_neigh.append(nn) + n_neigh = np.concatenate(n_neigh) + self._avg_num_neigh_approx = np.mean(n_neigh) + return self._avg_num_neigh_approx + + def run_stat(self, mode='stat'): + """ + Loop over dataset and init any statistics might need + Unlink SevenNetGraphDataset, neighbors count is not computed as + it requires to build graph + """ + if self._scanned is True: + return # statistics already computed + target_sequence = ( + self.stat_sequence if mode == 'stat' else np.arange(len(self)) + ) + y_keys: List[str] = [KEY.ENERGY, KEY.PER_ATOM_ENERGY, KEY.FORCE, KEY.STRESS] + natoms_counter = Counter() + composition = np.zeros((len(target_sequence), NUM_UNIV_ELEMENT)) + stats: Dict[str, Dict[str, Any]] = {y: {'_array': []} for y in y_keys} + + for i, atom_idx in enumerate( + tqdm(target_sequence, desc='run_stat', total=len(target_sequence)) + ): + atoms = self._dataset.get_atoms(atom_idx) + atoms = _set_atoms_y([atoms])[0] + z = atoms.get_atomic_numbers() + natoms_counter.update(z.tolist()) + composition[i] = np.bincount(z, minlength=NUM_UNIV_ELEMENT) + for y, dct in stats.items(): + if y == KEY.ENERGY: + dct['_array'].append(atoms.info['y_energy']) + elif y == KEY.PER_ATOM_ENERGY: + dct['_array'].append(atoms.info['y_energy'] / len(atoms)) + elif y == KEY.FORCE: + dct['_array'].append(atoms.arrays['y_force'].reshape(-1)) + elif y == KEY.STRESS: + dct['_array'].append(atoms.info['y_stress'].reshape(-1)) + + for y, dct in stats.items(): + if y == KEY.FORCE: + array = np.concatenate(dct['_array']) + else: + array = np.array(dct['_array']).reshape(-1) + dct.update( + { + 'mean': float(np.mean(array)), + 'std': float(np.std(array)), + 'median': float(np.quantile(array, q=0.5)), + 'max': float(np.max(array)), + 'min': float(np.min(array)), + '_array': array, + } + ) + + natoms = {chemical_symbols[int(z)]: cnt for z, cnt in natoms_counter.items()} + natoms['total'] = sum(list(natoms.values())) + self.statistics.update( + { + '_composition': composition, + '_natoms': natoms, + **stats, + } + ) + self._scanned = True + + +def _get_keys_from_config(config, start, end): + keys = [] + for k in config: + if k.startswith(start) and k.endswith(end): + keys.append(k) + return keys + + +def from_config( + config: dict[str, Any], + working_dir: str = os.getcwd(), + dataset_keys: Optional[list[str]] = None, + sequence_keys: Optional[list[str]] = None, +): + from sevenn.sevenn_logger import Logger + + log = Logger() + if dataset_keys is None: + dataset_keys = _get_keys_from_config(config, 'load_', '_path') + + if sequence_keys is None: + sequence_keys = _get_keys_from_config(config, 'load_', '_sequence') + + if KEY.LOAD_TRAINSET not in dataset_keys: + raise ValueError(f'{KEY.LOAD_TRAINSET} must be present in config') + + # initialize arguments for loading dataset + dataset_args = { + 'cutoff': config[KEY.CUTOFF], + **config[KEY.DATA_FORMAT_ARGS], + } + + chem_keys = [KEY.CHEMICAL_SPECIES, KEY.NUM_SPECIES, KEY.TYPE_MAP] + is_auto_mode = all([config[ck] == 'auto' for ck in chem_keys]) + + datasets: Dict[str, SevenNetASElmdbDataset] = {} + for dk in dataset_keys: + if not (paths := config[dk]): + continue + if isinstance(paths, str): + paths = [paths] + name = dk.split('_')[1].strip() + sk = dk.replace('_path', '_sequence') + + stat_sequence_info = 10000 + if sk in config: + stat_sequence_info = config[sk].get('stat_sequence_info', None) + dataset_args.update( + { + 'files': paths, + 'stat_sequence_info': stat_sequence_info, + 'is_auto_mode': is_auto_mode, + } + ) + datasets[name] = SevenNetASElmdbDataset(**dataset_args) + + if not config[KEY.COMPUTE_STATISTICS]: + log.writeline( + """ + Computing statistics is skipped, note that if any of other + configurations requires statistics (shift, scale, avg_num_neigh, + chemical_species as auto), SevenNet eventually raise an error! + """ + ) + return datasets + + train_set = datasets['trainset'] + chem_species = set(train_set.species) + + # print statistics of each dataset + for name, dataset in datasets.items(): + dataset.run_stat() + log.bar() + log.writeline(f'{name} distribution (may subsampled):') + log.statistic_write(dataset.statistics) + log.format_k_v('# atoms (node)', dataset.natoms, write=True) + # log.format_k_v('# structures (graph)', len(dataset), write=True) + log.format_k_v('# total structures in db', len(dataset._dataset), write=True) + + chem_species.update(dataset.species) + log.bar() + + # initialize known species from dataset if 'auto' + # sorted to alphabetical order (which is same as before) + if is_auto_mode: # see parse_input.py + log.writeline('Known species are obtained from the dataset') + config.update(util.chemical_species_preprocess(sorted(list(chem_species)))) + + # retrieve shift, scale, conv_denominaotrs from user input (keyword) + init_from_stats = [KEY.SHIFT, KEY.SCALE, KEY.CONV_DENOMINATOR] + for k in init_from_stats: + input = config[k] # statistic key or numbers + # If it is not 'str', 1: It is 'continue' training + # 2: User manually inserted numbers + if isinstance(input, str) and hasattr(train_set, input): + var = getattr(train_set, input) + config.update({k: var}) + log.writeline(f'{k} is obtained from statistics') + elif isinstance(input, str) and not hasattr(train_set, input): + raise NotImplementedError(input) + + return datasets diff --git a/sevenn/train/loss.py b/sevenn/train/loss.py index 3b271be7..941eabcd 100644 --- a/sevenn/train/loss.py +++ b/sevenn/train/loss.py @@ -2,6 +2,7 @@ import torch +import sevenn._const as CONST import sevenn._keys as KEY @@ -201,34 +202,157 @@ def _preprocess( return pred, ref, w_tensor +class L2Regularization(LossDefinition): + """ + L2 regularization for task-specific (modal) parameters. + Regularizes the last weight view of modal-specific IrrepsLinear layers, + which corresponds to the modal input dimension. + """ + + def __init__( + self, + name: str, + module_keys: List[str], + reg_modal_only: bool = True, + ): + super().__init__( + name=name, + unit=None, + criterion=None, + ref_key=None, + pred_key=None, + ) + self.module_keys = module_keys + self.reg_modal_only = reg_modal_only + + def get_loss(self, batch_data: Dict[str, Any], model: Optional[Callable] = None): + device = batch_data['x'].device + ret = torch.tensor([0.0], device=device) + for module_key in self.module_keys: + module = model._modules[module_key] # type: ignore + reg_params = list(module._modules['linear'].weight_views())[-1] + reg_loss = torch.sum(torch.pow(reg_params, 2)) + ret = ret + reg_loss + return 0.5 * ret # penalty = (1/2)||w||^2 + + def get_cosine( + self, batch_data: Dict[str, Any], model: Optional[Callable] = None + ): + cosine_list = [] + for module_key in self.module_keys: + module = model._modules[module_key] # type: ignore + reg_params = list(module._modules['linear'].weight_views())[-1] + dot = torch.dot(reg_params[0], reg_params[1]) + norm = torch.norm(reg_params[0]) * torch.norm(reg_params[1]) + cosine_list.append(dot / norm) + ret = torch.tensor( + [sum(cosine_list) / len(cosine_list)], + device=batch_data['x'].device, + ) + return ret + + +def get_modal_regularization( + config: Dict[str, Any], + model: Optional[torch.nn.Module] = None, +) -> Optional[Tuple[LossDefinition, float]]: + reg_params = config.get(KEY.REG_PARAM, {}) + + modal_param = reg_params.get('modal', {}) + if not modal_param or not config.get(KEY.USE_MODALITY, False): + return None + + if not model: + raise ValueError('modal reg is requested but model is not given.') + + module_keys_to_reg = [] + for module_key in list(model._modules.keys()): + for ( + use_modal_module_key, + modal_module_name, + ) in CONST.IMPLEMENTED_MODAL_MODULE_DICT.items(): + if ( + not config[use_modal_module_key] + or modal_module_name not in module_key + ): + continue + elif modal_module_name == 'reduce_input_to_hidden': + continue + module_keys_to_reg.append(module_key) + + return ( + L2Regularization('L2_modal', module_keys_to_reg, reg_modal_only=True), + float(modal_param.get(KEY.REG_WEIGHT, 1e-5)), + ) + + +def make_loss_info_dict_from_config(config: Dict[str, Any]): + # this is for backward compatibility + loss_info_dict = {} + loss_type = config.get(KEY.LOSS, 'mse').lower() + loss_param = config.get(KEY.LOSS_PARAM, {}) + for key in ['energy', 'force', 'stress']: + loss_info_dict[key] = {} + # loss_weight not initialized here. + loss_info_dict[key].update( + {KEY.LOSS_TYPE: loss_type, KEY.LOSS_PARAM: loss_param} + ) + + return loss_info_dict + + def get_loss_functions_from_config( config: Dict[str, Any], + model: Optional[torch.nn.Module] = None, ) -> List[Tuple[LossDefinition, float]]: from sevenn.train.optim import loss_dict + from sevenn.train.reewc.loss import get_ewc_loss loss_functions = [] # list of tuples (loss_definition, weight) - loss = loss_dict[config[KEY.LOSS].lower()] - loss_param = config.get(KEY.LOSS_PARAM, {}) + loss_info_dict = config.get(KEY.LOSS, 'mse') + if isinstance(loss_info_dict, str): + loss_info_dict = make_loss_info_dict_from_config(config) + + loss_function_cls_dict = { + 'energy': PerAtomEnergyLoss, + 'force': ForceLoss, + 'stress': StressLoss, + } + loss_weights = { + 'energy': config.get(KEY.ENERGY_WEIGHT, 1.0), + 'force': config[KEY.FORCE_WEIGHT], + 'stress': config[KEY.STRESS_WEIGHT], + } use_weight = config.get(KEY.USE_WEIGHT, False) - if use_weight: - loss_param['reduction'] = 'none' - criterion = loss(**loss_param) - commons = {'use_weight': use_weight} - loss_functions.append((PerAtomEnergyLoss(**commons), 1.0)) - loss_functions.append((ForceLoss(**commons), config[KEY.FORCE_WEIGHT])) + keys = ['energy', 'force'] if config[KEY.IS_TRAIN_STRESS]: - loss_functions.append((StressLoss(**commons), config[KEY.STRESS_WEIGHT])) - - for loss_function, _ in loss_functions: # why do these? - if loss_function.criterion is None: - loss_function.assign_criteria(criterion) - - from sevenn.train.reewc.loss import append_ewc_loss - - append_ewc_loss(loss_functions, config) + keys += ['stress'] + + for key in keys: + loss_info = loss_info_dict.get(key, {}) + loss_param = loss_info.get(KEY.LOSS_PARAM, {}) + loss_weight = loss_info.get(KEY.LOSS_WEIGHT, loss_weights[key]) + if (loss_type := loss_info.get(KEY.LOSS_TYPE, 'mse').lower()) == 'l2mae': + if key == 'energy': + raise NotImplementedError('L2MAE not implemented for energy.') + else: + loss_param.update({'prop': key}) + + loss_cls = loss_dict[loss_type] + if use_weight: + loss_param['reduction'] = 'none' + criterion = loss_cls(**loss_param) + loss_function_cls = loss_function_cls_dict[key] + loss_function = loss_function_cls(criterion=criterion, **commons) + loss_functions.append((loss_function, loss_weight)) + + if addi_loss := get_ewc_loss(config): + loss_functions.append(addi_loss) + if addi_loss := get_modal_regularization(config, model): + loss_functions.append(addi_loss) return loss_functions diff --git a/sevenn/train/optim.py b/sevenn/train/optim.py index 1a6c7d74..04f498db 100644 --- a/sevenn/train/optim.py +++ b/sevenn/train/optim.py @@ -6,6 +6,40 @@ from torch.optim import adagrad, adam, adamw, radam, sgd from torch.optim.lr_scheduler import _LRScheduler + +class L2MAE(nn.Module): + """ + L2 norm (Frobenius norm) based MAE loss. + For stress, norm of 3*3 matrix is used for invariance + """ + + def __init__( + self, + prop: str = 'force', + reduction: str = 'mean', + ): + super().__init__() + self.prop = prop + self.dim = 3 if prop == 'force' else 6 + self.reduction = reduction + + def forward(self, input, target): + if self.prop == 'force': + diff = input.view([-1, self.dim]) - target.view([-1, self.dim]) + else: + diff = input.view([-1, self.dim]) - target.view([-1, self.dim]) + diff = torch.cat((diff, diff[:, -3:]), dim=1) + norm = torch.norm(diff, p=2, dim=-1) + if self.reduction == 'none': + # make it (# component * atoms (or structures)) + # for consistency with MSE & MAE + return torch.repeat_interleave(norm, self.dim) + if self.reduction == 'mean': + return torch.mean(norm) + if self.reduction == 'sum': + return torch.sum(norm) + + optim_dict = { 'sgd': sgd.SGD, 'adagrad': adagrad.Adagrad, @@ -149,6 +183,12 @@ def step(self, epoch=None): 'cosineannealingwarmuplr': CosineAnnealingWarmupRestarts, 'reducelronplateau': scheduler.ReduceLROnPlateau, 'linearlr': scheduler.LinearLR, + 'onecyclelr': scheduler.OneCycleLR, } -loss_dict = {'mse': nn.MSELoss, 'huber': nn.HuberLoss} +loss_dict = { + 'mse': nn.MSELoss, + 'huber': nn.HuberLoss, + 'mae': nn.L1Loss, + 'l2mae': L2MAE, +} diff --git a/sevenn/train/reewc/__init__.py b/sevenn/train/reewc/__init__.py deleted file mode 100644 index d2640282..00000000 --- a/sevenn/train/reewc/__init__.py +++ /dev/null @@ -1,16 +0,0 @@ -from .loss import EWCLoss, append_ewc_loss -from .rehearsal import ( - build_memory_loader, - reewc_dataset_keys, - validate_reewc_config, -) -from .trainer import ReewcTrainer - -__all__ = [ - 'EWCLoss', - 'append_ewc_loss', - 'ReewcTrainer', - 'build_memory_loader', - 'reewc_dataset_keys', - 'validate_reewc_config', -] diff --git a/sevenn/train/reewc/loss.py b/sevenn/train/reewc/loss.py index e69d8933..0c7ffe5d 100644 --- a/sevenn/train/reewc/loss.py +++ b/sevenn/train/reewc/loss.py @@ -1,5 +1,5 @@ import warnings -from typing import Any, Callable, Dict, List, Optional, Tuple +from typing import Any, Callable, Dict, Optional, Tuple import torch @@ -57,9 +57,7 @@ def _check_and_align(self, model: Callable) -> None: f'{tuple(self.opt_params_dict[name].shape)}' ) - model_params = { - n: p for n, p in model.named_parameters() if p.requires_grad - } + model_params = {n: p for n, p in model.named_parameters() if p.requires_grad} if len(model_params) == 0: raise ValueError('EWC requires the model to have trainable parameters') @@ -92,9 +90,7 @@ def _check_and_align(self, model: Callable) -> None: self.to(next(iter(model_params.values())).device) self._checked = True - def get_loss( - self, batch_data: Dict[str, Any], model: Optional[Callable] = None - ): + def get_loss(self, batch_data: Dict[str, Any], model: Optional[Callable] = None): _ = batch_data if model is None: raise ValueError('EWCLoss requires the model to compute the penalty') @@ -113,10 +109,7 @@ def get_loss( return ewc_loss -def append_ewc_loss( - loss_functions: List[Tuple[LossDefinition, float]], - config: Dict[str, Any], -) -> None: +def get_ewc_loss(config: Dict[str, Any]) -> Optional[Tuple[LossDefinition, float]]: """reEWC: append the EWC penalty as an extra loss term when a precomputed Fisher information and reference parameters are given under continue.""" cont = config.get(KEY.CONTINUE, {}) @@ -124,7 +117,7 @@ def append_ewc_loss( opt_path = cont.get(KEY.OPT_PARAMS, False) ewc_lambda = cont.get(KEY.EWC_LAMBDA, 0) if not (fisher_path or opt_path or ewc_lambda): - return + return None if not (fisher_path and opt_path): raise ValueError( 'EWC requires both continue.fisher_information and ' @@ -138,4 +131,4 @@ def append_ewc_loss( raise ValueError('EWC requires continue.ewc_lambda > 0') fisher = torch.load(fisher_path, map_location='cpu', weights_only=True) opt = torch.load(opt_path, map_location='cpu', weights_only=True) - loss_functions.append((EWCLoss(fisher, opt), ewc_lambda / 2.0)) + return (EWCLoss(fisher, opt), ewc_lambda / 2.0) diff --git a/sevenn/train/reewc/trainer.py b/sevenn/train/reewc/trainer.py index 02cd7845..e5b21600 100644 --- a/sevenn/train/reewc/trainer.py +++ b/sevenn/train/reewc/trainer.py @@ -30,7 +30,7 @@ def from_config( ) -> 'ReewcTrainer': trainer = ReewcTrainer( model, - loss_functions=get_loss_functions_from_config(config), + loss_functions=get_loss_functions_from_config(config, model), optimizer_cls=optim_dict[config.get(KEY.OPTIMIZER, 'adam').lower()], optimizer_args=config.get(KEY.OPTIM_PARAM, {}), scheduler_cls=scheduler_dict[ diff --git a/sevenn/train/sampler.py b/sevenn/train/sampler.py new file mode 100644 index 00000000..9e5226f5 --- /dev/null +++ b/sevenn/train/sampler.py @@ -0,0 +1,95 @@ +import math +from typing import Iterator, List, Optional + +import numpy as np +import torch.utils.data.sampler + + +class OrderedSampler(torch.utils.data.sampler.Sampler): + """ + Deterministic sampler for DDP training with resume support. + Work both for single / multi GPU. + For single GPU, use world_size=1, rank=0 (default). + """ + + def __init__( + self, + dataset, + sequence: Optional[List[int]] = None, + shuffle: bool = False, + seed: int = 777, + world_size: int = 1, + rank: int = 0, + ): + if sequence is None: + self.sequence = np.arange(len(dataset)) + else: + self.sequence = np.array(sequence) + self.shuffle = shuffle + self.rng = np.random.default_rng(seed) + + assert world_size > 0 and rank < world_size + self.world_size = world_size + self.rank = rank + + self.total_samples_per_rank = math.ceil( + len(self.sequence) / self.world_size + ) + self.total_size = self.total_samples_per_rank * self.world_size + + self._start_index = 0 + + def continue_from_data_progress( + self, + numpy_rng_state: dict, + total_data_num: int = -1, + current_data_index: int = 0, + ): + if numpy_rng_state is not None: + self.rng.bit_generator.state = numpy_rng_state + if total_data_num < 0: # Nothing to continue + return + elif total_data_num != len(self.sequence): + raise ValueError( + 'data_progress not compatible' + + 'set reset_data_progress: True' + ) + self._start_index = current_data_index + + def get_rng_state(self): + return self.rng.bit_generator.state + + def permutate_sequence(self): + return self.rng.permutation(self.sequence) + + def refresh_sequence(self): + self._start_index = 0 + + def __iter__(self) -> Iterator[int]: + indices = self.sequence.copy() + if self.shuffle: + # deterministically shuffle based on numpy rng state + indices = self.permutate_sequence() + + # add extra samples to make it evenly divisible + padding_size = self.total_size - len(indices) + if padding_size <= len(indices): + padding_sequence = indices[:padding_size] + else: + rep_num = math.ceil(padding_size / len(indices)) + padding_sequence = np.tile(indices, rep_num)[:padding_size] + indices = np.concatenate((indices, padding_sequence)) + assert len(indices) == self.total_size + + # subsample + indices = indices[ + self._start_index + self.rank : self.total_size : self.world_size + ] + assert len(indices) == len(self) + self.refresh_sequence() # after one epoch, it initializes to 0. + + return iter(indices) + + def __len__(self) -> int: + current_idx_per_rank = int(self._start_index / self.world_size) + return self.total_samples_per_rank - current_idx_per_rank diff --git a/sevenn/train/trainer.py b/sevenn/train/trainer.py index 22d952e8..44fe16a3 100644 --- a/sevenn/train/trainer.py +++ b/sevenn/train/trainer.py @@ -38,10 +38,11 @@ def __init__( self, model: torch.nn.Module, loss_functions: List[Tuple[LossDefinition, float]], - optimizer_cls, + optimizer_cls=None, optimizer_args: Optional[Dict[str, Any]] = None, scheduler_cls=None, scheduler_args: Optional[Dict[str, Any]] = None, + grad_clip_norm_th: Optional[float] = None, device: Union[torch.device, str] = 'auto', distributed: bool = False, distributed_backend: str = 'nccl', @@ -80,6 +81,7 @@ def __init__( else: self.scheduler = None self.loss_functions = loss_functions + self.grad_clip_norm_th = grad_clip_norm_th @staticmethod def from_config( @@ -88,13 +90,14 @@ def from_config( ) -> 'Trainer': trainer = Trainer( model, - loss_functions=get_loss_functions_from_config(config), + loss_functions=get_loss_functions_from_config(config, model), 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, {}), + grad_clip_norm_th=config.get(KEY.GRAD_CLIP, None), device=config.get(KEY.DEVICE, 'auto'), distributed=config.get(KEY.IS_DDP, False), distributed_backend=config.get(KEY.DDP_BACKEND, 'nccl'), @@ -120,7 +123,7 @@ def args_from_checkpoint(checkpoint: str) -> Tuple[Dict, Dict, Dict]: 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) + loss_functions = get_loss_functions_from_config(config, model) return ( { @@ -146,7 +149,7 @@ def run_one_epoch( """ Run single epoch with given dataloader Args: - loader: iterable yieds AtomGraphData + loader: iterable yields 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 @@ -159,6 +162,7 @@ def run_one_epoch( if wrap_tqdm: total_len = wrap_tqdm if isinstance(wrap_tqdm, int) else None loader = tqdm(loader, total=total_len) + _model = self.model if not self.distributed else self.model.module for _, batch in enumerate(loader): if is_train: @@ -166,19 +170,64 @@ def run_one_epoch( batch = batch.to(self.device, non_blocking=True) output = self.model(batch) if error_recorder is not None: - error_recorder.update(output) + error_recorder.update(output, model=_model) 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) + indv_loss = loss_def.get_loss(output, _model) if indv_loss is not None: - total_loss += (indv_loss * w) + total_loss += indv_loss * w total_loss.backward() - self.optimizer.step() + if self.grad_clip_norm_th is not None: + torch.nn.utils.clip_grad_norm_( + self.model.parameters(), + max_norm=self.grad_clip_norm_th, + ) + self.optimizer.step() # DDP syncs weight here if self.distributed and error_recorder is not None: self.recorder_all_reduce(error_recorder) + def train_one_batch( + self, + batch, + error_recorder: Optional[ErrorRecorder] = None, + ) -> None: + """ + Train on a single batch. Used for batch-level training loop. + + Args: + batch: AtomGraphData batch + error_recorder: ErrorRecorder instance to compute errors + """ + self.model.train() + self.optimizer.zero_grad() + batch = batch.to(self.device, non_blocking=True) + output = self.model(batch) + _model = self.model if not self.distributed else self.model.module + + if error_recorder is not None: + error_recorder.update(output, model=_model) + + total_loss = torch.tensor([0.0], device=self.device) + for loss_def, w in self.loss_functions: + total_loss += loss_def.get_loss(output, _model) * w + + total_loss.backward() + + # Grad clipping + if self.grad_clip_norm_th is not None: + norm = torch.nn.utils.clip_grad_norm_( + self.model.parameters(), + max_norm=self.grad_clip_norm_th, + ) + if norm > self.grad_clip_norm_th and self.rank == 0: + print( + f'[Clipping] Grad norm {norm:.2f} into {self.grad_clip_norm_th}', + flush=True, + ) + self.optimizer.step() + def scheduler_step(self, metric: Optional[float] = None) -> None: if self.scheduler is None: return @@ -201,6 +250,7 @@ def get_checkpoint_dict(self) -> Dict[str, Any]: 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(), diff --git a/sevenn/util.py b/sevenn/util.py index 55d7b24d..89d54065 100644 --- a/sevenn/util.py +++ b/sevenn/util.py @@ -3,7 +3,7 @@ import pathlib import shutil import warnings -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union +from typing import TYPE_CHECKING, Any, Dict, List, Tuple, Union import numpy as np import requests diff --git a/tests/unit_tests/test_batch_d3.py b/tests/unit_tests/test_batch_d3.py index 48da698e..139a2996 100644 --- a/tests/unit_tests/test_batch_d3.py +++ b/tests/unit_tests/test_batch_d3.py @@ -2,7 +2,6 @@ # TODO: check stress-things with non-pbc input import numpy as np import pytest -from ase import Atoms from ase.build import bulk, molecule try: diff --git a/tests/unit_tests/test_train.py b/tests/unit_tests/test_train.py index a7d831c9..ab36db56 100644 --- a/tests/unit_tests/test_train.py +++ b/tests/unit_tests/test_train.py @@ -166,7 +166,7 @@ def test_processing_continue_v2_7net0(tmp_path): conv_denominator_ref = np.array([35.989574] * 5) with Logger().switch_file(str(tmp_path / 'log.sevenn')): - state_dicts, epoch = processing_continue_v2(cfg) + state_dicts, epoch, _ = processing_continue_v2(cfg) assert epoch == 601 assert np.allclose(np.array(cfg['shift']), shift_ref) assert np.allclose(np.array(cfg['shift'])[0], -5.062768)