From 2530a3ae710b434e94fa52f598810d44418a2396 Mon Sep 17 00:00:00 2001 From: alphalm4 Date: Thu, 29 Jan 2026 08:42:10 +0900 Subject: [PATCH 01/32] add l2mae, ordered sampler, batch training --- CHANGELOG.md | 7 +- sevenn/_const.py | 10 +- sevenn/_keys.py | 9 ++ sevenn/scripts/processing_by_batch.py | 205 ++++++++++++++++++++++++++ sevenn/scripts/processing_continue.py | 38 ++++- sevenn/scripts/train.py | 161 ++++++++++++++++++-- sevenn/train/optim.py | 43 +++++- sevenn/train/sampler.py | 97 ++++++++++++ sevenn/train/trainer.py | 67 ++++++++- 9 files changed, 612 insertions(+), 25 deletions(-) create mode 100644 sevenn/scripts/processing_by_batch.py create mode 100644 sevenn/train/sampler.py diff --git a/CHANGELOG.md b/CHANGELOG.md index d7ff2120..58795577 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,11 @@ All notable changes to this project will be documented in this file. ## [0.12.1] +### Added +- SevenNet-Omni-i8, SevenNet-Omni-i12 +- L2MAE loss +- OrderedSampler, batch training + ### Fixed - FlashTP with LAMMPS parallel in torch @@ -10,8 +15,6 @@ All notable changes to this project will be documented in this file. ### 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/sevenn/_const.py b/sevenn/_const.py index 0004e68c..66297cf0 100644 --- a/sevenn/_const.py +++ b/sevenn/_const.py @@ -254,15 +254,19 @@ 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.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, }, @@ -286,16 +290,20 @@ 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.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 2f430b1b..1a9a192b 100644 --- a/sevenn/_keys.py +++ b/sevenn/_keys.py @@ -116,8 +116,11 @@ 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' @@ -130,6 +133,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' @@ -149,6 +153,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' diff --git a/sevenn/scripts/processing_by_batch.py b/sevenn/scripts/processing_by_batch.py new file mode 100644 index 00000000..bcf06830 --- /dev/null +++ b/sevenn/scripts/processing_by_batch.py @@ -0,0 +1,205 @@ +import os +import time +from copy import deepcopy +from typing import 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 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', '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..2e4d173e 100644 --- a/sevenn/scripts/processing_continue.py +++ b/sevenn/scripts/processing_continue.py @@ -13,15 +13,19 @@ ) -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: + For epoch training: (state_dicts, epoch) + For batch training: (state_dicts, epoch, data_progress) """ 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,16 +83,38 @@ 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, ] + + # Handle data progress for batch training + 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]}') + #log.writeline(f'Checkpoint previous epoch was: {from_epoch}') # duplicated? + + log.writeline('checkpoint loading successful') + return state_dicts, epoch, data_progress + + log.writeline('checkpoint loading successful') return state_dicts, epoch diff --git a/sevenn/scripts/train.py b/sevenn/scripts/train.py index 28371857..5eb33d13 100644 --- a/sevenn/scripts/train.py +++ b/sevenn/scripts/train.py @@ -1,5 +1,7 @@ +import math 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 @@ -15,33 +17,134 @@ 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 + """ + is_train = dataset_key == 'trainset' batch_size = config[KEY.BATCH_SIZE] + + if isinstance(dataset, dict): + batch_size = dataset.get('batch_size', batch_size) + dataset = dataset['dataset'] + shuffle = is_train and config[KEY.TRAIN_SHUFFLE] + train_by_batch = config.get(KEY.TRAIN_BY_BATCH, False) 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 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 + else: + world_size, rank = 1, 0 + + # Use OrderedSampler for batch training mode to preserve data order + if train_by_batch and is_train: + from sevenn.train.sampler import OrderedSampler + + seed = config.get(KEY.RANDOM_SEED, None) + try: + sequence = config[f'load_{dataset_key}_sequence']['total_sequence_path'] + except: + sequence = None + if sequence is not None: # when using custom sequence (e.g. subset) + sequence = np.load(sequence) + sampler = OrderedSampler(dataset, sequence, shuffle, seed, world_size, rank) + loader_args.update({'sampler': sampler}) + loader_args.pop( + 'shuffle', None + ) # sampler is mutually exclusive with shuffle + return DataLoader(**loader_args) +def update_config_for_batch_training( + config: Dict[str, Any], + train_loader + ) -> None: + """ + Update scheduler parameters for batch-level training. + + This converts epoch-based scheduler parameters to step-based parameters + when using batch training mode. + """ + if not config.get(KEY.TRAIN_BY_BATCH, False): + return + + # 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, + ) + ) + + +# 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.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 from .processing_continue import processing_continue_v2 from .processing_epoch import processing_epoch_v2 @@ -55,25 +158,47 @@ def train_v2(config: Dict[str, Any], working_dir: str) -> None: log.writeline('***************************************************') config[KEY.LOAD_TRAINSET] = config.pop(KEY.LOAD_DATASET) + # Initialize data progress for batch training + train_by_batch = config.get(KEY.TRAIN_BY_BATCH, False) + if train_by_batch: + data_progress = { + KEY.TOTAL_DATA_NUM: -1, + KEY.CURRENT_DATA_IDX: 0, + KEY.NUMPY_RNG_STATE: None, + } + else: + data_progress = {} # dummy + # 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) + result = processing_continue_v2(config) + if train_by_batch: + state_dicts, start_epoch, data_progress = result + else: + state_dicts, start_epoch = result + # Load datasets based on type + dataset_type = config[KEY.DATASET_TYPE] if config.get(KEY.USE_MODALITY, False): 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) - elif config[KEY.DATASET_TYPE] == 'atoms': + elif dataset_type == 'atoms': datasets = atoms_dataset.from_config(config, working_dir) 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')) + k: loader_from_config(config, v, dataset_key=k) for k, v in datasets.items() } + # Update scheduler config for batch training + if train_by_batch: + update_config_for_batch_training(config, loaders['trainset']) + log.write('\nModel building...\n') model = build_E3_equivariant_model(config) log.print_model_info(model, config) @@ -82,9 +207,19 @@ def train_v2(config: Dict[str, Any], working_dir: str) -> None: 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') @@ -110,7 +245,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')) + k: loader_from_config(config, v, dataset_key=k) for k, v in datasets.items() } loaders = list(loaders.values()) diff --git a/sevenn/train/optim.py b/sevenn/train/optim.py index 10e75790..68afe0c3 100644 --- a/sevenn/train/optim.py +++ b/sevenn/train/optim.py @@ -1,7 +1,42 @@ +import torch import torch.nn as nn import torch.optim.lr_scheduler as scheduler from torch.optim import adagrad, adam, adamw, radam, sgd + +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, @@ -18,6 +53,12 @@ 'cosineannealinglr': scheduler.CosineAnnealingLR, '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/sampler.py b/sevenn/train/sampler.py new file mode 100644 index 00000000..2cc71a89 --- /dev/null +++ b/sevenn/train/sampler.py @@ -0,0 +1,97 @@ +import math +from typing import Iterator, List, Optional + +import numpy as np +import torch.utils.data.sampler + +from torch_geometric.data import Dataset + + +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 9a7e186e..4f4ae8ab 100644 --- a/sevenn/train/trainer.py +++ b/sevenn/train/trainer.py @@ -42,6 +42,7 @@ def __init__( 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(model: torch.nn.Module, config: Dict[str, Any]) -> 'Trainer': @@ -92,6 +94,7 @@ def from_config(model: torch.nn.Module, config: Dict[str, Any]) -> 'Trainer': 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'), @@ -142,7 +145,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 @@ -169,11 +172,70 @@ def run_one_epoch( if indv_loss is not None: 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) + + 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() + + # TODO: NaN sanitizer - replace NaN/Inf gradients with zero + # for name, p in self.model.named_parameters(): + # if p.grad is None: + # continue + # if not torch.isfinite(p.grad).all(): + # if self.rank == 0: + # print( + # f'[nan2zero] NaN/Inf gradient detected in {p.shape}, ' + # 'resetting to 0', + # flush=True, + # ) + # p.grad = torch.nan_to_num(p.grad, nan=0.0, posinf=0.0, neginf=0.0) + + # 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 ' + f'{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 @@ -196,6 +258,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(), From 51c1b2b05eb3cbe0d15ab92aa0bdfeb5e8368c3f Mon Sep 17 00:00:00 2001 From: alphalm4 Date: Sun, 1 Feb 2026 14:15:42 +0900 Subject: [PATCH 02/32] add train_shift and train_scale --- sevenn/_const.py | 4 ++++ sevenn/_keys.py | 2 ++ sevenn/model_build.py | 17 +++++++++++++---- sevenn/nn/scale.py | 27 +++++++++++++++++++++------ sevenn/scripts/processing_continue.py | 7 ++++++- 5 files changed, 46 insertions(+), 11 deletions(-) diff --git a/sevenn/_const.py b/sevenn/_const.py index 66297cf0..22258fb8 100644 --- a/sevenn/_const.py +++ b/sevenn/_const.py @@ -116,6 +116,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, @@ -157,6 +159,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, diff --git a/sevenn/_keys.py b/sevenn/_keys.py index 1a9a192b..7d124eca 100644 --- a/sevenn/_keys.py +++ b/sevenn/_keys.py @@ -225,6 +225,8 @@ 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 diff --git a/sevenn/model_build.py b/sevenn/model_build.py index 3117ce29..e9a20346 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/scripts/processing_continue.py b/sevenn/scripts/processing_continue.py index 2e4d173e..c1331d42 100644 --- a/sevenn/scripts/processing_continue.py +++ b/sevenn/scripts/processing_continue.py @@ -151,7 +151,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 From a50958f13827a5f7a10d99b15070068754c88936 Mon Sep 17 00:00:00 2001 From: alphalm4 Date: Mon, 2 Feb 2026 01:43:04 +0900 Subject: [PATCH 03/32] l2reg; tmp --- sevenn/_const.py | 11 +++ sevenn/_keys.py | 3 + sevenn/error_recorder.py | 116 +++++++++++++++++++---- sevenn/scripts/convert_model_modality.py | 8 +- sevenn/scripts/processing_by_batch.py | 2 +- sevenn/scripts/processing_epoch.py | 2 +- sevenn/train/loss.py | 96 +++++++++++++++++++ sevenn/train/trainer.py | 29 ++++-- 8 files changed, 233 insertions(+), 34 deletions(-) diff --git a/sevenn/_const.py b/sevenn/_const.py index 22258fb8..e08db04f 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'] @@ -262,6 +271,7 @@ def data_defaults(config): 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, @@ -298,6 +308,7 @@ def data_defaults(config): 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: lambda x: type(x) in [float, int], diff --git a/sevenn/_keys.py b/sevenn/_keys.py index 7d124eca..475b3b75 100644 --- a/sevenn/_keys.py +++ b/sevenn/_keys.py @@ -234,6 +234,9 @@ USE_FLASH_TP = 'use_flash_tp' CUEQUIVARIANCE_CONFIG = 'cuequivariance_config' +REG_PARAM = 'regularization_param' +REG_WEIGHT = 'regularization_weight' + _NORMALIZE_SPH = '_normalize_sph' OPTIMIZE_BY_REDUCE = 'optimize_by_reduce' diff --git a/sevenn/error_recorder.py b/sevenn/error_recorder.py index 262ea06f..cbf6ae62 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) @@ -127,7 +145,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 +200,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 +226,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 +252,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 +274,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 +300,39 @@ 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 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 +344,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 +381,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} @@ -408,11 +466,21 @@ def init_total_loss_metric( def from_config( config: Dict[str, Any], loss_functions: Optional[List[Tuple[LossDefinition, float]]] = None, + reg_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 + if loss_functions is not None: + all_loss_functions = ( + loss_functions + reg_functions + if isinstance(reg_functions, list) + else loss_functions + ) + else: + all_loss_functions = None + err_config = config.get(KEY.ERROR_RECORD, False) if not err_config: raise ValueError( @@ -432,17 +500,25 @@ def from_config( if err_type == 'TotalLoss': # special case err_metrics.append( ErrorRecorder.init_total_loss_metric( - config, criteria, loss_functions + config, criteria, all_loss_functions ) ) continue + elif err_type == 'Modal_cos': # special case + metric_cls = ModalWeightCosine + metric_kwargs['loss_def'], _ = _get_loss_function_from_name( + all_loss_functions, 'L2_modal' + ) + metric_kwargs.pop('unit', None) + err_metrics.append(metric_cls(**metric_kwargs)) + continue metric_cls = ErrorRecorder.METRIC_DICT[metric_name] assert isinstance(metric_kwargs['name'], str) if metric_name == 'Loss': - if loss_functions is not None: + if all_loss_functions is not None: metric_cls = LossError metric_kwargs['loss_def'], _ = _get_loss_function_from_name( - loss_functions, metric_kwargs['name'] + all_loss_functions, metric_kwargs['name'] ) else: metric_cls = CustomError 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 index bcf06830..3de00513 100644 --- a/sevenn/scripts/processing_by_batch.py +++ b/sevenn/scripts/processing_by_batch.py @@ -59,7 +59,7 @@ def processing_by_batch( 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 + config, trainer.loss_functions, trainer.reg_functions ) recorders = {k: deepcopy(recorder) for k in loaders} diff --git a/sevenn/scripts/processing_epoch.py b/sevenn/scripts/processing_epoch.py index 38e671db..4de088e7 100644 --- a/sevenn/scripts/processing_epoch.py +++ b/sevenn/scripts/processing_epoch.py @@ -38,7 +38,7 @@ def processing_epoch_v2( best_metric = best_metric or config.get(KEY.BEST_METRIC, 'TotalLoss') assert isinstance(best_metric, str) recorder = error_recorder or ErrorRecorder.from_config( - config, trainer.loss_functions + config, trainer.loss_functions, trainer.reg_functions ) recorders = {k: deepcopy(recorder) for k in loaders} diff --git a/sevenn/train/loss.py b/sevenn/train/loss.py index a6f8a769..0bc80c55 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,6 +202,101 @@ 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 ret + + 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_module_keys_for_reg( + config: Dict[str, Any], all_module_keys: List[str] +) -> List[str]: + module_keys_to_reg = [] + for module_key in all_module_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 module_keys_to_reg + + +def get_regularization_from_config( + config: Dict[str, Any], all_module_keys: List[str] +) -> List[Tuple[LossDefinition, float]]: + reg_params = config.get(KEY.REG_PARAM, {}) + reg_functions: List[Tuple[LossDefinition, float]] = [] + + modal_param = reg_params.get('modal', {}) + if not modal_param: + return reg_functions + + reg_weight = float(modal_param.get(KEY.REG_WEIGHT, 1e-5)) + module_keys_to_reg = _get_modal_module_keys_for_reg( + config, all_module_keys + ) + + reg_functions.append(( + L2Regularization('L2_modal', module_keys_to_reg, reg_modal_only=True), + reg_weight, + )) + + return reg_functions + + def get_loss_functions_from_config( config: Dict[str, Any], ) -> List[Tuple[LossDefinition, float]]: diff --git a/sevenn/train/trainer.py b/sevenn/train/trainer.py index 4f4ae8ab..51ca0686 100644 --- a/sevenn/train/trainer.py +++ b/sevenn/train/trainer.py @@ -11,9 +11,9 @@ import sevenn._keys as KEY from sevenn.error_recorder import ErrorRecorder -from sevenn.train.loss import LossDefinition +from sevenn.train.loss import L2Regularization, LossDefinition -from .loss import get_loss_functions_from_config +from .loss import get_loss_functions_from_config, get_regularization_from_config from .optim import optim_dict, scheduler_dict @@ -38,7 +38,8 @@ def __init__( self, model: torch.nn.Module, loss_functions: List[Tuple[LossDefinition, float]], - optimizer_cls, + reg_functions: Optional[List[Tuple[L2Regularization, float]]] = None, + optimizer_cls=None, optimizer_args: Optional[Dict[str, Any]] = None, scheduler_cls=None, scheduler_args: Optional[Dict[str, Any]] = None, @@ -81,13 +82,18 @@ def __init__( else: self.scheduler = None self.loss_functions = loss_functions + self.reg_functions = reg_functions or [] self.grad_clip_norm_th = grad_clip_norm_th @staticmethod def from_config(model: torch.nn.Module, config: Dict[str, Any]) -> 'Trainer': + reg_functions = get_regularization_from_config( + config, list(model._modules.keys()) + ) trainer = Trainer( model, loss_functions=get_loss_functions_from_config(config), + reg_functions=reg_functions, optimizer_cls=optim_dict[config.get(KEY.OPTIMIZER, 'adam').lower()], optimizer_args=config.get(KEY.OPTIM_PARAM, {}), scheduler_cls=scheduler_dict[ @@ -121,11 +127,15 @@ def args_from_checkpoint(checkpoint: str) -> Tuple[Dict, Dict, Dict]: optimizer_cls = optim_dict[config[KEY.OPTIMIZER].lower()] scheduler_cls = scheduler_dict[config[KEY.SCHEDULER].lower()] loss_functions = get_loss_functions_from_config(config) + reg_functions = get_regularization_from_config( + config, list(model._modules.keys()) + ) return ( { 'model': model, 'loss_functions': loss_functions, + 'reg_functions': reg_functions, 'optimizer_cls': optimizer_cls, 'optimizer_args': config[KEY.OPTIM_PARAM], 'scheduler_cls': scheduler_cls, @@ -158,19 +168,23 @@ 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: 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) + 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) + for reg_def, w in self.reg_functions: + reg_loss = reg_def.get_loss(output, _model) + total_loss += reg_loss * w / 2 total_loss.backward() if self.grad_clip_norm_th is not None: torch.nn.utils.clip_grad_norm_( @@ -201,11 +215,14 @@ def train_one_batch( _model = self.model if not self.distributed else self.model.module if error_recorder is not None: - error_recorder.update(output) + 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 + for reg_def, w in self.reg_functions: + reg_loss = reg_def.get_loss(output, _model) + total_loss += reg_loss * w / 2 total_loss.backward() From 532159845cddd38f307b188cb5f9e301d6518fe9 Mon Sep 17 00:00:00 2001 From: Jaesun0912 Date: Fri, 6 Feb 2026 19:49:44 +0900 Subject: [PATCH 04/32] add aselmdb --- pyproject.toml | 4 +- setup.cfg | 2 +- sevenn/_const.py | 3 +- sevenn/_keys.py | 3 + sevenn/checkpoint.py | 12 + sevenn/error_recorder.py | 14 +- sevenn/scripts/processing_epoch.py | 2 +- sevenn/scripts/train.py | 31 +- sevenn/train/aselmdb_dataset.py | 815 +++++++++++++++++++++++++++++ sevenn/train/loss.py | 65 ++- 10 files changed, 929 insertions(+), 22 deletions(-) create mode 100644 sevenn/train/aselmdb_dataset.py diff --git a/pyproject.toml b/pyproject.toml index 733fa34f..43717fb2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,7 +30,9 @@ dependencies = [ "pandas", "requests", "ninja", - "setuptools>=61.0" + "setuptools>=61.0", + "lmdb", + "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 e08db04f..afc2b16c 100644 --- a/sevenn/_const.py +++ b/sevenn/_const.py @@ -223,6 +223,7 @@ def model_defaults(config): KEY.USE_MODAL_WISE_SCALE: False, KEY.SHIFT: 'per_atom_energy_mean', KEY.SCALE: 'force_rms', + KEY.LOADER_KWARGS: {}, # KEY.DATA_SHUFFLE: True, # KEY.DATA_WEIGHT: False, # KEY.DATA_MODALITY: False, @@ -238,7 +239,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'], # 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, diff --git a/sevenn/_keys.py b/sevenn/_keys.py index 475b3b75..fe0a0a88 100644 --- a/sevenn/_keys.py +++ b/sevenn/_keys.py @@ -112,6 +112,8 @@ EPOCH = 'epoch' LOSS = 'loss' LOSS_PARAM = 'loss_param' +LOSS_TYPE = 'loss_type' +LOSS_WEIGHT = 'loss_weight' OPTIMIZER = 'optimizer' OPTIM_PARAM = 'optim_param' SCHEDULER = 'scheduler' @@ -219,6 +221,7 @@ 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' diff --git a/sevenn/checkpoint.py b/sevenn/checkpoint.py index ec95cffa..e289a914 100644 --- a/sevenn/checkpoint.py +++ b/sevenn/checkpoint.py @@ -191,6 +191,11 @@ 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 + reset_optimizer: False + reset_scheduler: False + reset_data_progress: False + reset_epoch: False self._model_state_dict = None self._optimizer_state_dict = None self._scheduler_state_dict = None @@ -269,6 +274,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: @@ -293,6 +304,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 cbf6ae62..7acf456a 100644 --- a/sevenn/error_recorder.py +++ b/sevenn/error_recorder.py @@ -468,9 +468,16 @@ def from_config( loss_functions: Optional[List[Tuple[LossDefinition, float]]] = None, reg_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_info_dict = config[KEY.LOSS] + if isinstance(loss_info_dict, str): + loss_info_dict = make_loss_info_dict_from_config(config) + + criteria_dict = {} + for err_type in ['Energy' ,'Force', 'Stress']: + loss_cls = loss_dict[loss_info_dict.get(KEY.LOSS_TYPE, 'mse').lower()] + loss_param = loss_info_dict.get(KEY.LOSS_PARAM, {}) + criteria = loss_cls(**loss_param) if loss_functions is None else None + criteria_dict[err_type] = criteria if loss_functions is not None: all_loss_functions = ( @@ -497,6 +504,7 @@ def from_config( err_metrics = [] for err_type, metric_name in err_config: metric_kwargs = get_err_type(err_type) + criteria = criteria_dict.get(err_type, None) if err_type == 'TotalLoss': # special case err_metrics.append( ErrorRecorder.init_total_loss_metric( diff --git a/sevenn/scripts/processing_epoch.py b/sevenn/scripts/processing_epoch.py index 4de088e7..4f94a855 100644 --- a/sevenn/scripts/processing_epoch.py +++ b/sevenn/scripts/processing_epoch.py @@ -62,7 +62,7 @@ def processing_epoch_v2( f.write(','.join(head) + '\n') if start_epoch == 1: - path = f'{prefix}/checkpoint_0.pth' # save first epoch + path = f'{prefix}/checkpoint_initial.pth' # save first epoch trainer.write_checkpoint(path, config=config, epoch=0) for epoch in range(start_epoch, total_epoch + 1): # one indexing diff --git a/sevenn/scripts/train.py b/sevenn/scripts/train.py index 5eb33d13..2089ec55 100644 --- a/sevenn/scripts/train.py +++ b/sevenn/scripts/train.py @@ -45,6 +45,9 @@ def loader_from_config( 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) + if config[KEY.IS_DDP]: dist.barrier() world_size = dist.get_world_size() @@ -131,6 +134,25 @@ def update_config_for_batch_training( ) +def datasets_from_py(config, script): + import importlib.util + from pathlib import Path + + 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) + spec.loader.exec_module(module) + + 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: """ @@ -181,12 +203,19 @@ def train_v2(config: Dict[str, Any], working_dir: str) -> None: # Load datasets based on type dataset_type = config[KEY.DATASET_TYPE] - if config.get(KEY.USE_MODALITY, False): + if ( + config.get(KEY.USE_MODALITY, False) + and not config[KEY.DATASET_TYPE] == 'custom' + ): datasets = modal_dataset.from_config(config, working_dir) elif dataset_type == 'graph': datasets = graph_dataset.from_config(config, working_dir) 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: {dataset_type}') diff --git a/sevenn/train/aselmdb_dataset.py b/sevenn/train/aselmdb_dataset.py new file mode 100644 index 00000000..d4a1a2d3 --- /dev/null +++ b/sevenn/train/aselmdb_dataset.py @@ -0,0 +1,815 @@ +from __future__ import annotations + +import os +import os.path as osp +from glob import glob +import time +import warnings +from collections import Counter +from pathlib import Path +import bisect +import zlib +import typing +from typing import Any, Callable, Dict, List, Optional, Union + +import numpy as np +import torch.distributed as dist +import lmdb +import orjson + +import ase +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 + # inefficient for large dataset. If the db we're using already presents a list of + # ids and there is no query, we can just use that list instead and save ourselves + # 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 sequence: + self.total_sequence = np.array(sequence) + else: + _seq = list(range(len(self._dataset))) + np.random.shuffle(_seq) + self.total_sequence = _seq + """ + + 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.total_sequence), stat_sequence_info) + sample_num = min(len(self), stat_sequence_info) + """ + self.stat_sequence = np.random.choice( + self.total_sequence, sample_num, replace=False + ) + """ + 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.total_sequence) * stat_sequence_info) + sample_num = int(len(self) * stat_sequence_info) + """ + self.stat_sequence = np.random.choice( + self.total_sequence, sample_num, replace=False + ) + """ + 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)}' + ) + #self._run_sequence = self.total_sequence + + 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 + + """ + def continue_from_data_progress( + self, + total_data_num: int = -1, + current_data_index: int = 0, + sequence: Optional[List[int]] = None, + ): + if total_data_num < 0: # Nothing to continue + return + elif total_data_num != len(self._dataset): + raise ValueError( + 'data_progress is not compatible with the dataset' + + 'set reset_data_progress: True to fresh start' + ) + if sequence is not None: + assert len(sequence) == len(self._dataset) + self.total_sequence = sequence + self._truncated_sequence(current_data_index) + + def set_epoch(self, epoch: int, is_ddp: bool): + ''' + Should be called before every epoch + Mimic behavior of distributed sampler + ''' + # TODO: rngkey things based on epoch? + self.shuffle_sequence(is_ddp) + self._run_sequence = self.total_sequence + + def _truncated_sequence(self, index): + ''' + Used to start from middle of index (for continuing large-data training) + Also changes __len__ of the dataset + ''' + self._run_sequence = self.total_sequence[index:] + + def shuffle_sequence(self, broadcast=False): + # ambiguous as both run_sequence and total_sequece can be random idx + # assume total_sequnce is the one that is shuffled every epoch and + # run_sequence is simlpy for truncation of total_sequence for continue + shuffled = np.random.permutation(self.total_sequence) + if broadcast: + shuffled_bcast = [shuffled] + dist.broadcast_object_list(shuffled_bcast, src=0) + shuffled = shuffled_bcast[0] + self.total_sequence = shuffled + + def save_sequence(self, filename): # should not used + np.save(filename, self.total_sequence) + """ + @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 self.total_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') + + total_sequence_path = None + stat_sequence_info = 10000 + if sk in config: + #total_sequence_path = config[sk].get('total_sequence_path', None) + stat_sequence_info = config[sk].get('stat_sequence_info', None) + dataset_args.update( + { + 'files': paths, + #'sequence_file': total_sequence_path, + '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 0bc80c55..fa27611a 100644 --- a/sevenn/train/loss.py +++ b/sevenn/train/loss.py @@ -297,30 +297,67 @@ def get_regularization_from_config( return reg_functions +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], + config: Dict[str, Any] ) -> List[Tuple[LossDefinition, float]]: from sevenn.train.optim import loss_dict 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) + 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)) return loss_functions + From 9025cce7a76762757e830b1f81de4d80e003f9fc Mon Sep 17 00:00:00 2001 From: Jaesun0912 Date: Fri, 6 Feb 2026 21:22:46 +0900 Subject: [PATCH 05/32] fix error_recorder bug --- sevenn/error_recorder.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sevenn/error_recorder.py b/sevenn/error_recorder.py index 7acf456a..6be01cd3 100644 --- a/sevenn/error_recorder.py +++ b/sevenn/error_recorder.py @@ -14,7 +14,7 @@ import torch.distributed as dist import sevenn._keys as KEY -from sevenn.train.loss import LossDefinition +from sevenn.train.loss import LossDefinition, make_loss_info_dict_from_config from .train.optim import loss_dict From c6a5ce30bdfd8b117b53478c70ea1348eda51587 Mon Sep 17 00:00:00 2001 From: alphalm4 Date: Fri, 27 Feb 2026 19:04:58 +0900 Subject: [PATCH 06/32] fix --- sevenn/_const.py | 2 +- sevenn/scripts/train.py | 9 +++++---- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/sevenn/_const.py b/sevenn/_const.py index afc2b16c..99466269 100644 --- a/sevenn/_const.py +++ b/sevenn/_const.py @@ -239,7 +239,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', 'custom'], + 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, diff --git a/sevenn/scripts/train.py b/sevenn/scripts/train.py index 2089ec55..cdb1c417 100644 --- a/sevenn/scripts/train.py +++ b/sevenn/scripts/train.py @@ -72,16 +72,16 @@ def loader_from_config( sampler = OrderedSampler(dataset, sequence, shuffle, seed, world_size, rank) loader_args.update({'sampler': sampler}) loader_args.pop( - 'shuffle', None + 'shuffle', None ) # sampler is mutually exclusive with shuffle return DataLoader(**loader_args) def update_config_for_batch_training( - config: Dict[str, Any], - train_loader - ) -> None: + config: Dict[str, Any], + train_loader +) -> None: """ Update scheduler parameters for batch-level training. @@ -162,6 +162,7 @@ def train_v2(config: Dict[str, Any], working_dir: str) -> None: - 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 3c665d9da03e5e3f7264659adb5d2eac81d7ea2f Mon Sep 17 00:00:00 2001 From: alphalm4 Date: Sun, 15 Mar 2026 19:30:25 +0900 Subject: [PATCH 07/32] rebase changelog --- CHANGELOG.md | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7c25f4ae..a58fe212 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,17 +4,14 @@ All notable changes to this project will be documented in this file. ## [0.12.2.dev] ### Added - Support OpenEquivariance +- L2MAE loss +- OrderedSampler, batch training ### Changed - **[Breaking]** Rename optional dependency group `mliap` into `mliap12` (reflecting its CUDA 12.x dependency). - Add `cueq13` and `mliap13` optional dependency groups for CUDA 13.x. ## [0.12.1] -### Added -- SevenNet-Omni-i8, SevenNet-Omni-i12 -- L2MAE loss -- OrderedSampler, batch training - ### Fixed - FlashTP with LAMMPS parallel in torch - Single-atom inference failure with ASE+flashTP, LAMMPS-Torch and LAMMPS ML-IAP @@ -27,6 +24,7 @@ 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 From 81cce2ab609578d4b939e5e3c414dfd6b207c5b2 Mon Sep 17 00:00:00 2001 From: alphalm4 Date: Thu, 28 May 2026 16:10:30 +0900 Subject: [PATCH 08/32] lmdb ver less than 2 required for simultaneous open --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index bfc6f77a..46faadf0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,7 +31,7 @@ dependencies = [ "requests", "ninja", "setuptools>=61.0", - "lmdb", + "lmdb<2.0.0", "orjson" ] [project.optional-dependencies] From 04052f22852f4378c9d31a66e2a385596b7013be Mon Sep 17 00:00:00 2001 From: alphalm4 Date: Thu, 28 May 2026 16:10:57 +0900 Subject: [PATCH 09/32] port load_validset_sequence --- sevenn/parse_input.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/sevenn/parse_input.py b/sevenn/parse_input.py index 099e4c03..5960642c 100644 --- a/sevenn/parse_input.py +++ b/sevenn/parse_input.py @@ -186,6 +186,10 @@ 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'): + data_meta[k] = config[k] # TODO: check validity in here, not `from_config` + for key, default in _const.DEFAULT_DATA_CONFIG.items(): data_meta[key] = config_initialize( key, config, default, _const.DATA_CONFIG_CONDITION From bbd043b78fdaa76002b328d753bce4c6b3e21288 Mon Sep 17 00:00:00 2001 From: alphalm4 Date: Thu, 28 May 2026 16:13:06 +0900 Subject: [PATCH 10/32] remove deprecated lines in aselmdb_dataset.py --- sevenn/train/aselmdb_dataset.py | 74 --------------------------------- 1 file changed, 74 deletions(-) diff --git a/sevenn/train/aselmdb_dataset.py b/sevenn/train/aselmdb_dataset.py index d4a1a2d3..9727152b 100644 --- a/sevenn/train/aselmdb_dataset.py +++ b/sevenn/train/aselmdb_dataset.py @@ -481,38 +481,18 @@ def __init__( self.statistics = {} self._dataset = AseDBDataset(src=files) - """ - if sequence: - self.total_sequence = np.array(sequence) - else: - _seq = list(range(len(self._dataset))) - np.random.shuffle(_seq) - self.total_sequence = _seq - """ 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.total_sequence), stat_sequence_info) sample_num = min(len(self), stat_sequence_info) - """ - self.stat_sequence = np.random.choice( - self.total_sequence, sample_num, replace=False - ) - """ 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.total_sequence) * stat_sequence_info) sample_num = int(len(self) * stat_sequence_info) - """ - self.stat_sequence = np.random.choice( - self.total_sequence, sample_num, replace=False - ) - """ self.stat_sequence = np.random.choice( np.arange(len(self)), sample_num, replace=False ) @@ -522,7 +502,6 @@ def __init__( 'stat_sequence_info should be one of str, int, float, ' + f'but got {type(stat_sequence_info)}' ) - #self._run_sequence = self.total_sequence def __len__(self): # total, run_sequence deprecated. @@ -564,55 +543,6 @@ def set_atoms_y_with_idx(self, idx): atoms = _set_atoms_y([atoms])[0] return atoms - """ - def continue_from_data_progress( - self, - total_data_num: int = -1, - current_data_index: int = 0, - sequence: Optional[List[int]] = None, - ): - if total_data_num < 0: # Nothing to continue - return - elif total_data_num != len(self._dataset): - raise ValueError( - 'data_progress is not compatible with the dataset' - + 'set reset_data_progress: True to fresh start' - ) - if sequence is not None: - assert len(sequence) == len(self._dataset) - self.total_sequence = sequence - self._truncated_sequence(current_data_index) - - def set_epoch(self, epoch: int, is_ddp: bool): - ''' - Should be called before every epoch - Mimic behavior of distributed sampler - ''' - # TODO: rngkey things based on epoch? - self.shuffle_sequence(is_ddp) - self._run_sequence = self.total_sequence - - def _truncated_sequence(self, index): - ''' - Used to start from middle of index (for continuing large-data training) - Also changes __len__ of the dataset - ''' - self._run_sequence = self.total_sequence[index:] - - def shuffle_sequence(self, broadcast=False): - # ambiguous as both run_sequence and total_sequece can be random idx - # assume total_sequnce is the one that is shuffled every epoch and - # run_sequence is simlpy for truncation of total_sequence for continue - shuffled = np.random.permutation(self.total_sequence) - if broadcast: - shuffled_bcast = [shuffled] - dist.broadcast_object_list(shuffled_bcast, src=0) - shuffled = shuffled_bcast[0] - self.total_sequence = shuffled - - def save_sequence(self, filename): # should not used - np.save(filename, self.total_sequence) - """ @property def species(self): mode = ( @@ -654,7 +584,6 @@ def run_stat(self, mode='stat'): if self._scanned is True: return # statistics already computed target_sequence = ( - #self.stat_sequence if mode == 'stat' else self.total_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] @@ -752,15 +681,12 @@ def from_config( name = dk.split('_')[1].strip() sk = dk.replace('_path', '_sequence') - total_sequence_path = None stat_sequence_info = 10000 if sk in config: - #total_sequence_path = config[sk].get('total_sequence_path', None) stat_sequence_info = config[sk].get('stat_sequence_info', None) dataset_args.update( { 'files': paths, - #'sequence_file': total_sequence_path, 'stat_sequence_info': stat_sequence_info, 'is_auto_mode': is_auto_mode, } From 1c83a3cb2dbb61b2427c19f4604f44bc6b558550 Mon Sep 17 00:00:00 2001 From: alphalm4 Date: Thu, 28 May 2026 22:08:38 +0900 Subject: [PATCH 11/32] fix validset subsampling --- sevenn/scripts/train.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/sevenn/scripts/train.py b/sevenn/scripts/train.py index cdb1c417..991b66c5 100644 --- a/sevenn/scripts/train.py +++ b/sevenn/scripts/train.py @@ -59,7 +59,10 @@ def loader_from_config( world_size, rank = 1, 0 # Use OrderedSampler for batch training mode to preserve data order - if train_by_batch and is_train: + # verified only for validset + # TODO: I think 'train_by_batch' and 'sampling validset' is independent, + # so 'train_by_batch' should be removed + if train_by_batch: from sevenn.train.sampler import OrderedSampler seed = config.get(KEY.RANDOM_SEED, None) From 39a48c9931abdcd988490c53ea87a3085a3eb37d Mon Sep 17 00:00:00 2001 From: alphalm4 Date: Sat, 30 May 2026 17:05:49 +0900 Subject: [PATCH 12/32] restore error_recorder (loss dict is build from loss.py) --- sevenn/error_recorder.py | 30 ++++++++++++++++-------------- 1 file changed, 16 insertions(+), 14 deletions(-) diff --git a/sevenn/error_recorder.py b/sevenn/error_recorder.py index 6be01cd3..f646459c 100644 --- a/sevenn/error_recorder.py +++ b/sevenn/error_recorder.py @@ -14,7 +14,7 @@ import torch.distributed as dist import sevenn._keys as KEY -from sevenn.train.loss import LossDefinition, make_loss_info_dict_from_config +from sevenn.train.loss import LossDefinition from .train.optim import loss_dict @@ -121,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__( @@ -468,16 +469,18 @@ def from_config( loss_functions: Optional[List[Tuple[LossDefinition, float]]] = None, reg_functions: Optional[List[Tuple[LossDefinition, float]]] = None, ) -> 'ErrorRecorder': - loss_info_dict = config[KEY.LOSS] - if isinstance(loss_info_dict, str): - loss_info_dict = make_loss_info_dict_from_config(config) - - criteria_dict = {} - for err_type in ['Energy' ,'Force', 'Stress']: - loss_cls = loss_dict[loss_info_dict.get(KEY.LOSS_TYPE, 'mse').lower()] - loss_param = loss_info_dict.get(KEY.LOSS_PARAM, {}) - criteria = loss_cls(**loss_param) if loss_functions is None else None - criteria_dict[err_type] = criteria + 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 if loss_functions is not None: all_loss_functions = ( @@ -504,7 +507,6 @@ def from_config( err_metrics = [] for err_type, metric_name in err_config: metric_kwargs = get_err_type(err_type) - criteria = criteria_dict.get(err_type, None) if err_type == 'TotalLoss': # special case err_metrics.append( ErrorRecorder.init_total_loss_metric( From 8bd596dea12bee2b347e6a60c76a62c932bf7808 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 30 May 2026 08:17:01 +0000 Subject: [PATCH 13/32] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- sevenn/train/aselmdb_dataset.py | 76 ++++++++++++++++----------------- sevenn/train/loss.py | 1 - sevenn/train/sampler.py | 1 - 3 files changed, 37 insertions(+), 41 deletions(-) diff --git a/sevenn/train/aselmdb_dataset.py b/sevenn/train/aselmdb_dataset.py index 9727152b..4a4379a7 100644 --- a/sevenn/train/aselmdb_dataset.py +++ b/sevenn/train/aselmdb_dataset.py @@ -1,23 +1,22 @@ from __future__ import annotations +import bisect import os import os.path as osp -from glob import glob import time +import typing import warnings +import zlib from collections import Counter +from glob import glob from pathlib import Path -import bisect -import zlib -import typing from typing import Any, Callable, Dict, List, Optional, Union -import numpy as np -import torch.distributed as dist +import ase import lmdb +import numpy as np import orjson - -import ase +import torch.distributed as dist from ase.data import chemical_symbols from ase.db.core import Database, now, ops from ase.db.row import AtomsRow @@ -31,7 +30,6 @@ from sevenn.train.dataload import _set_atoms_y - class LMDBDatabase(Database): """ Same class develop by Meta @@ -129,46 +127,46 @@ def _write( else: row = AtomsRow(atoms) row.ctime = mtime - row.user = os.getenv("USER") + row.user = os.getenv('USER') dct = {} for key in row.__dict__: - if key[0] == "_" or key in row._keys or key == "id": + if key[0] == '_' or key in row._keys or key == 'id': continue dct[key] = row[key] - dct["mtime"] = mtime + dct['mtime'] = mtime if key_value_pairs: - dct["key_value_pairs"] = key_value_pairs + dct['key_value_pairs'] = key_value_pairs if data: - dct["data"] = data + dct['data'] = data - constraints = row.get("constraints") + constraints = row.get('constraints') if constraints: - dct["constraints"] = [constraint.todict() for constraint in 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"]) + 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")) + data = self.txn.get(f"{idx}".encode('ascii')) assert data is not None # Add the new entry self.txn.put( - f"{idx}".encode("ascii"), + 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"), + 'nextid'.encode('ascii'), zlib.compress(orjson.dumps(nextid, option=orjson.OPT_SERIALIZE_NUMPY)), ) # check if id is in removed ids and remove accordingly @@ -191,7 +189,7 @@ def _update( def _write_deleted_ids(self): self.txn.put( - "deleted_ids".encode("ascii"), + 'deleted_ids'.encode('ascii'), zlib.compress( orjson.dumps(self.deleted_ids, option=orjson.OPT_SERIALIZE_NUMPY) ), @@ -199,7 +197,7 @@ def _write_deleted_ids(self): def delete(self, ids: list[int]) -> None: for idx in ids: - self.txn.delete(f"{idx}".encode("ascii")) + self.txn.delete(f"{idx}".encode('ascii')) self.ids.remove(idx) self.deleted_ids += ids @@ -209,7 +207,7 @@ 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")) + data = self.txn.get(f"{idx}".encode('ascii')) if data is not None: dct = orjson.loads(zlib.decompress(data)) @@ -217,14 +215,14 @@ def _get_row(self, idx: int, include_data: bool = True): raise KeyError(f"Id {idx} missing from the database!") if not include_data: - dct.pop("data", None) + dct.pop('data', None) - dct["id"] = idx + 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")) + data = self.txn.get(f"{self.ids[index]}".encode('ascii')) if data is not None: dct = orjson.loads(zlib.decompress(data)) @@ -232,9 +230,9 @@ def _get_row_by_index(self, index: int, include_data: bool = True): raise KeyError(f"Id {id} missing from the database!") if not include_data: - dct.pop("data", None) + dct.pop('data', None) - dct["id"] = id + dct['id'] = id return AtomsRow(dct) def _select( @@ -247,14 +245,14 @@ def _select( offset: int = 0, sort: str | None = None, include_data: bool = True, - columns: str = "all", + columns: str = 'all', ): if explain: - yield {"explain": (0, 0, 0, "scan table")} + yield {'explain': (0, 0, 0, 'scan table')} return if sort is not None: - if sort[0] == "-": + if sort[0] == '-': reverse = True sort = sort[1:] else: @@ -297,9 +295,9 @@ def _select( 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 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: @@ -311,7 +309,7 @@ def _select( def metadata(self): """Load the metadata from the DB if present""" if self._metadata is None: - metadata = self.txn.get("metadata".encode("ascii")) + metadata = self.txn.get('metadata'.encode('ascii')) if metadata is None: self._metadata = {} else: @@ -325,7 +323,7 @@ def metadata(self, dct): # Put the updated metadata dictionary self.txn.put( - "metadata".encode("ascii"), + 'metadata'.encode('ascii'), zlib.compress(orjson.dumps(dct, option=orjson.OPT_SERIALIZE_NUMPY)), ) @@ -333,7 +331,7 @@ def metadata(self, dct): def _nextid(self): """Get the id of the next row to be written""" # Get the nextid - nextid_data = self.txn.get("nextid".encode("ascii")) + 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: @@ -361,7 +359,7 @@ def _load_ids(self) -> None: """ # Load the deleted ids - deleted_ids_data = self.txn.get("deleted_ids".encode("ascii")) + 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)) @@ -409,7 +407,7 @@ def __init__( # a lot of time! self.db_ids = [] for db in self.dbs: - if hasattr(db, "ids"): + if hasattr(db, 'ids'): self.db_ids.append(db.ids) else: # this is the slow alternative diff --git a/sevenn/train/loss.py b/sevenn/train/loss.py index fa27611a..05cd42c2 100644 --- a/sevenn/train/loss.py +++ b/sevenn/train/loss.py @@ -360,4 +360,3 @@ def get_loss_functions_from_config( loss_functions.append((loss_function, loss_weight)) return loss_functions - diff --git a/sevenn/train/sampler.py b/sevenn/train/sampler.py index 2cc71a89..fa27435d 100644 --- a/sevenn/train/sampler.py +++ b/sevenn/train/sampler.py @@ -3,7 +3,6 @@ import numpy as np import torch.utils.data.sampler - from torch_geometric.data import Dataset From cdca4ac20d4f224b425ed0c1090d53a31eff7eb8 Mon Sep 17 00:00:00 2001 From: alphalm4 Date: Sat, 30 May 2026 17:36:58 +0900 Subject: [PATCH 14/32] lint --- sevenn/checkpoint.py | 4 ---- sevenn/parse_input.py | 3 ++- sevenn/scripts/processing_by_batch.py | 2 +- sevenn/scripts/processing_continue.py | 4 ++-- sevenn/train/aselmdb_dataset.py | 30 ++++++++++++++------------- 5 files changed, 21 insertions(+), 22 deletions(-) diff --git a/sevenn/checkpoint.py b/sevenn/checkpoint.py index 019e8202..6dbfbb5b 100644 --- a/sevenn/checkpoint.py +++ b/sevenn/checkpoint.py @@ -192,10 +192,6 @@ def __init__(self, checkpoint_path: Union[pathlib.Path, str]) -> None: self._config = None self._epoch = None self._data_progress = None - reset_optimizer: False - reset_scheduler: False - reset_data_progress: False - reset_epoch: False self._model_state_dict = None self._optimizer_state_dict = None self._scheduler_state_dict = None diff --git a/sevenn/parse_input.py b/sevenn/parse_input.py index 5960642c..d302431b 100644 --- a/sevenn/parse_input.py +++ b/sevenn/parse_input.py @@ -188,7 +188,8 @@ def init_data_config(config: Dict[str, Any]) -> Dict[str, Any]: for k in config: if k.startswith('load_') and k.endswith('_sequence'): - data_meta[k] = config[k] # TODO: check validity in here, not `from_config` + # 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( diff --git a/sevenn/scripts/processing_by_batch.py b/sevenn/scripts/processing_by_batch.py index 3de00513..79be8772 100644 --- a/sevenn/scripts/processing_by_batch.py +++ b/sevenn/scripts/processing_by_batch.py @@ -112,7 +112,7 @@ def processing_by_batch( ) scheduler_update_every_batch = ( - config.get(KEY.SCHEDULER_BATCH_MODE, False) + config.get(KEY.SCHEDULER_BATCH_MODE, False) ) # TODO: too long, refactor more diff --git a/sevenn/scripts/processing_continue.py b/sevenn/scripts/processing_continue.py index c1331d42..3a1de27a 100644 --- a/sevenn/scripts/processing_continue.py +++ b/sevenn/scripts/processing_continue.py @@ -108,8 +108,8 @@ def processing_continue_v2(config: Dict[str, Any]): 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]}') - #log.writeline(f'Checkpoint previous epoch was: {from_epoch}') # duplicated? + log.writeline(f'data index start from {data_progress[KEY.CURRENT_DATA_IDX]}') # noqa: E501 + # log.writeline(f'Checkpoint previous epoch was: {from_epoch}') # duplicated? # noqa: E501 log.writeline('checkpoint loading successful') return state_dicts, epoch, data_progress diff --git a/sevenn/train/aselmdb_dataset.py b/sevenn/train/aselmdb_dataset.py index 4a4379a7..f875d73f 100644 --- a/sevenn/train/aselmdb_dataset.py +++ b/sevenn/train/aselmdb_dataset.py @@ -167,7 +167,9 @@ def _write( self.ids.append(idx) self.txn.put( 'nextid'.encode('ascii'), - zlib.compress(orjson.dumps(nextid, option=orjson.OPT_SERIALIZE_NUMPY)), + 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: @@ -185,7 +187,9 @@ def _update( # 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) + self._write( + atoms=row, idx=idx, key_value_pairs=key_value_pairs, data=data + ) def _write_deleted_ids(self): self.txn.put( @@ -364,7 +368,9 @@ def _load_ids(self) -> 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)] + self.ids = [ + i for i in range(1, self._nextid) if i not in set(self.deleted_ids) + ] class AseDBDataset: @@ -399,11 +405,10 @@ def __init__( 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 - # inefficient for large dataset. If the db we're using already presents a list of - # ids and there is no query, we can just use that list instead and save ourselves + # 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: @@ -422,18 +427,15 @@ def __init__( 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) @@ -456,7 +458,7 @@ def __init__( self, cutoff: float, files: Union[str, List[str]], - #sequence: Optional[List[int]] = None, + # 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 @@ -504,12 +506,12 @@ def __init__( def __len__(self): # total, run_sequence deprecated. # Should only be used in OrderedSampler.__init__ - #return len(self._run_sequence) + # 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) + # 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) From eff657bcad6fa73a870f372c28f8b376d5cbc0b8 Mon Sep 17 00:00:00 2001 From: YutackPark Date: Tue, 9 Jun 2026 12:42:49 +0900 Subject: [PATCH 15/32] continue refactor --- sevenn/scripts/processing_continue.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/sevenn/scripts/processing_continue.py b/sevenn/scripts/processing_continue.py index 3a1de27a..6e578949 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 @@ -111,11 +111,11 @@ def processing_continue_v2(config: Dict[str, Any]): log.writeline(f'data index start from {data_progress[KEY.CURRENT_DATA_IDX]}') # noqa: E501 # log.writeline(f'Checkpoint previous epoch was: {from_epoch}') # duplicated? # noqa: E501 - log.writeline('checkpoint loading successful') + log.writeline('checkpoint loading success') return state_dicts, epoch, data_progress - - log.writeline('checkpoint loading successful') - return state_dicts, epoch + else: + log.writeline('checkpoint loading success') + return state_dicts, epoch def check_config_compatible(config: Dict[str, Any], config_cp: Dict[str, Any]): From 01341be26fd67313a8d8b096d4cb02020001256f Mon Sep 17 00:00:00 2001 From: YutackPark Date: Tue, 9 Jun 2026 12:56:43 +0900 Subject: [PATCH 16/32] refactor train.py --- sevenn/scripts/processing_continue.py | 7 ++- sevenn/scripts/train.py | 71 +++++++++++---------------- 2 files changed, 31 insertions(+), 47 deletions(-) diff --git a/sevenn/scripts/processing_continue.py b/sevenn/scripts/processing_continue.py index 6e578949..01e44a6c 100644 --- a/sevenn/scripts/processing_continue.py +++ b/sevenn/scripts/processing_continue.py @@ -98,6 +98,7 @@ def processing_continue_v2(config: Dict[str, Any]): ] # Handle data progress for batch training + data_progress = {} if train_by_batch: data_progress = { KEY.TOTAL_DATA_NUM: -1, @@ -112,10 +113,8 @@ def processing_continue_v2(config: Dict[str, Any]): # log.writeline(f'Checkpoint previous epoch was: {from_epoch}') # duplicated? # noqa: E501 log.writeline('checkpoint loading success') - return state_dicts, epoch, data_progress - else: - log.writeline('checkpoint loading success') - return state_dicts, epoch + + return state_dicts, epoch, data_progress or {} def check_config_compatible(config: Dict[str, Any], config_cp: Dict[str, Any]): diff --git a/sevenn/scripts/train.py b/sevenn/scripts/train.py index 991b66c5..144ef0eb 100644 --- a/sevenn/scripts/train.py +++ b/sevenn/scripts/train.py @@ -1,4 +1,6 @@ +import importlib.util import math +from pathlib import Path from typing import Any, Dict, List, Optional import numpy as np @@ -13,6 +15,7 @@ from sevenn.scripts.processing_continue import ( convert_modality_of_checkpoint_state_dct, ) +from sevenn.train.sampler import OrderedSampler from sevenn.train.trainer import Trainer @@ -30,15 +33,13 @@ def loader_from_config( (or dict with 'dataset' and 'batch_size') dataset_key: Key identifying the dataset """ - is_train = dataset_key == 'trainset' batch_size = config[KEY.BATCH_SIZE] if isinstance(dataset, dict): batch_size = dataset.get('batch_size', batch_size) dataset = dataset['dataset'] - shuffle = is_train and config[KEY.TRAIN_SHUFFLE] - train_by_batch = config.get(KEY.TRAIN_BY_BATCH, False) + shuffle = config[KEY.TRAIN_SHUFFLE] sampler = None loader_args = {'dataset': dataset, 'batch_size': batch_size, 'shuffle': shuffle} @@ -48,6 +49,7 @@ def loader_from_config( 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() world_size = dist.get_world_size() @@ -55,45 +57,36 @@ def loader_from_config( sampler = DistributedSampler(dataset, world_size, rank, shuffle=shuffle) loader_args.update({'sampler': sampler}) loader_args.pop('shuffle') # sampler is mutually exclusive with shuffle - else: - world_size, rank = 1, 0 # 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 train_by_batch: - from sevenn.train.sampler import OrderedSampler - - seed = config.get(KEY.RANDOM_SEED, None) - try: - sequence = config[f'load_{dataset_key}_sequence']['total_sequence_path'] - except: - sequence = None - if sequence is not None: # when using custom sequence (e.g. subset) - sequence = np.load(sequence) - sampler = OrderedSampler(dataset, sequence, shuffle, seed, world_size, rank) - loader_args.update({'sampler': sampler}) - loader_args.pop( - 'shuffle', None - ) # sampler is mutually exclusive with shuffle + if config.get(KEY.TRAIN_BY_BATCH, False): + sequence = config[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 update_config_for_batch_training( - config: Dict[str, Any], - train_loader -) -> None: +def update_config_for_batch_training(config: Dict[str, Any], train_loader) -> None: """ Update scheduler parameters for batch-level training. This converts epoch-based scheduler parameters to step-based parameters when using batch training mode. """ - if not config.get(KEY.TRAIN_BY_BATCH, False): - return - # 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( @@ -138,9 +131,6 @@ def update_config_for_batch_training( def datasets_from_py(config, script): - import importlib.util - from pathlib import Path - if isinstance(script, list): assert len(script) == 1, 'Need single python script' script = script[0] @@ -148,8 +138,8 @@ def datasets_from_py(config, script): 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) - spec.loader.exec_module(module) + 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 @@ -186,24 +176,21 @@ def train_v2(config: Dict[str, Any], working_dir: str) -> None: # 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, } - else: - data_progress = {} # dummy # config updated start_epoch = 1 state_dicts: Optional[List[dict]] = None if config[KEY.CONTINUE][KEY.CHECKPOINT]: - result = processing_continue_v2(config) - if train_by_batch: - state_dicts, start_epoch, data_progress = result - else: - state_dicts, start_epoch = result + # data_progress is non-empty only if train_by_batch is True + state_dicts, start_epoch, data_progress = processing_continue_v2(config) # Load datasets based on type dataset_type = config[KEY.DATASET_TYPE] @@ -224,8 +211,7 @@ def train_v2(config: Dict[str, Any], working_dir: str) -> None: raise ValueError(f'Unknown dataset type: {dataset_type}') loaders = { - k: loader_from_config(config, v, dataset_key=k) - for k, v in datasets.items() + k: loader_from_config(config, v, dataset_key=k) for k, v in datasets.items() } # Update scheduler config for batch training @@ -278,8 +264,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, dataset_key=k) - 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()) From cbe70d1c8aab3483c756026ac533a218a59c0d42 Mon Sep 17 00:00:00 2001 From: alphalm4 Date: Tue, 9 Jun 2026 13:06:00 +0900 Subject: [PATCH 17/32] fix --- sevenn/scripts/processing_epoch.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sevenn/scripts/processing_epoch.py b/sevenn/scripts/processing_epoch.py index 4f94a855..4de088e7 100644 --- a/sevenn/scripts/processing_epoch.py +++ b/sevenn/scripts/processing_epoch.py @@ -62,7 +62,7 @@ def processing_epoch_v2( f.write(','.join(head) + '\n') if start_epoch == 1: - path = f'{prefix}/checkpoint_initial.pth' # save first epoch + path = f'{prefix}/checkpoint_0.pth' # save first epoch trainer.write_checkpoint(path, config=config, epoch=0) for epoch in range(start_epoch, total_epoch + 1): # one indexing From 9a41f0add4add9435ebd5c37426c1b1b3abeeaa0 Mon Sep 17 00:00:00 2001 From: YutackPark Date: Tue, 9 Jun 2026 13:06:13 +0900 Subject: [PATCH 18/32] refactor loss.py --- sevenn/train/loss.py | 30 +++++++++++------------------- sevenn/train/trainer.py | 13 ------------- 2 files changed, 11 insertions(+), 32 deletions(-) diff --git a/sevenn/train/loss.py b/sevenn/train/loss.py index 05cd42c2..bf174018 100644 --- a/sevenn/train/loss.py +++ b/sevenn/train/loss.py @@ -254,9 +254,18 @@ def get_cosine( return ret -def _get_modal_module_keys_for_reg( +def get_regularization_from_config( config: Dict[str, Any], all_module_keys: List[str] -) -> List[str]: +) -> List[Tuple[LossDefinition, float]]: + reg_params = config.get(KEY.REG_PARAM, {}) + reg_functions: List[Tuple[LossDefinition, float]] = [] + + modal_param = reg_params.get('modal', {}) + if not modal_param: + return reg_functions + + reg_weight = float(modal_param.get(KEY.REG_WEIGHT, 1e-5)) + module_keys_to_reg = [] for module_key in all_module_keys: for ( @@ -271,23 +280,6 @@ def _get_modal_module_keys_for_reg( elif modal_module_name == 'reduce_input_to_hidden': continue module_keys_to_reg.append(module_key) - return module_keys_to_reg - - -def get_regularization_from_config( - config: Dict[str, Any], all_module_keys: List[str] -) -> List[Tuple[LossDefinition, float]]: - reg_params = config.get(KEY.REG_PARAM, {}) - reg_functions: List[Tuple[LossDefinition, float]] = [] - - modal_param = reg_params.get('modal', {}) - if not modal_param: - return reg_functions - - reg_weight = float(modal_param.get(KEY.REG_WEIGHT, 1e-5)) - module_keys_to_reg = _get_modal_module_keys_for_reg( - config, all_module_keys - ) reg_functions.append(( L2Regularization('L2_modal', module_keys_to_reg, reg_modal_only=True), diff --git a/sevenn/train/trainer.py b/sevenn/train/trainer.py index 51ca0686..ff3895fd 100644 --- a/sevenn/train/trainer.py +++ b/sevenn/train/trainer.py @@ -226,19 +226,6 @@ def train_one_batch( total_loss.backward() - # TODO: NaN sanitizer - replace NaN/Inf gradients with zero - # for name, p in self.model.named_parameters(): - # if p.grad is None: - # continue - # if not torch.isfinite(p.grad).all(): - # if self.rank == 0: - # print( - # f'[nan2zero] NaN/Inf gradient detected in {p.shape}, ' - # 'resetting to 0', - # flush=True, - # ) - # p.grad = torch.nan_to_num(p.grad, nan=0.0, posinf=0.0, neginf=0.0) - # Grad clipping if self.grad_clip_norm_th is not None: norm = torch.nn.utils.clip_grad_norm_( From 15f4767ea3e1803a3de47b8ee281f72bcef82d8d Mon Sep 17 00:00:00 2001 From: YutackPark Date: Tue, 9 Jun 2026 13:12:09 +0900 Subject: [PATCH 19/32] bugfix unittest --- tests/unit_tests/test_train.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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) From 295124f859ae455568e68fb23bf48aa107250be3 Mon Sep 17 00:00:00 2001 From: alphalm4 Date: Mon, 22 Jun 2026 16:45:25 +0900 Subject: [PATCH 20/32] add fallback in sequence loader --- sevenn/scripts/train.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sevenn/scripts/train.py b/sevenn/scripts/train.py index 144ef0eb..22cd5942 100644 --- a/sevenn/scripts/train.py +++ b/sevenn/scripts/train.py @@ -63,7 +63,7 @@ def loader_from_config( # 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[f'load_{dataset_key}_sequence'].get( + sequence = config.get(f'load_{dataset_key}_sequence', {}).get( 'total_sequence_path', None ) From 3761771bc666c9daad5536e785f110547892df75 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 22 Jun 2026 07:53:50 +0000 Subject: [PATCH 21/32] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- sevenn/scripts/train.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sevenn/scripts/train.py b/sevenn/scripts/train.py index 8b2f515e..9fdc99be 100644 --- a/sevenn/scripts/train.py +++ b/sevenn/scripts/train.py @@ -15,13 +15,13 @@ from sevenn.scripts.processing_continue import ( convert_modality_of_checkpoint_state_dct, ) -from sevenn.train.sampler import OrderedSampler from sevenn.train.reewc import ( ReewcTrainer, build_memory_loader, reewc_dataset_keys, validate_reewc_config, ) +from sevenn.train.sampler import OrderedSampler from sevenn.train.trainer import Trainer From 788a0eb0dce1ea0b09a5dac4cf24eb950caed8ea Mon Sep 17 00:00:00 2001 From: alphalm4 Date: Mon, 22 Jun 2026 16:57:18 +0900 Subject: [PATCH 22/32] lint --- sevenn/scripts/train.py | 2 +- sevenn/train/sampler.py | 1 - sevenn/util.py | 2 +- tests/unit_tests/test_batch_d3.py | 1 - 4 files changed, 2 insertions(+), 4 deletions(-) diff --git a/sevenn/scripts/train.py b/sevenn/scripts/train.py index 9fdc99be..2158555f 100644 --- a/sevenn/scripts/train.py +++ b/sevenn/scripts/train.py @@ -208,7 +208,7 @@ def train_v2(config: Dict[str, Any], working_dir: str) -> None: datasets = modal_dataset.from_config(config, working_dir) elif dataset_type == 'graph': datasets = graph_dataset.from_config( - config, working_dir, dataset_keys=reewc_dataset_keys(config) + config, working_dir, dataset_keys=reewc_dataset_keys(config) ) elif dataset_type == 'atoms': datasets = atoms_dataset.from_config(config, working_dir) diff --git a/sevenn/train/sampler.py b/sevenn/train/sampler.py index fa27435d..9e5226f5 100644 --- a/sevenn/train/sampler.py +++ b/sevenn/train/sampler.py @@ -3,7 +3,6 @@ import numpy as np import torch.utils.data.sampler -from torch_geometric.data import Dataset class OrderedSampler(torch.utils.data.sampler.Sampler): 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: From d8f2139bd94f2df0ee3dd921c505494afeace2b8 Mon Sep 17 00:00:00 2001 From: alphalm4 Date: Mon, 22 Jun 2026 17:21:00 +0900 Subject: [PATCH 23/32] append modal L2 reg into get_loss_functions_from_config --- sevenn/train/loss.py | 11 ++++++++++- sevenn/train/trainer.py | 24 ++++++------------------ 2 files changed, 16 insertions(+), 19 deletions(-) diff --git a/sevenn/train/loss.py b/sevenn/train/loss.py index 6066ecdf..c66553f8 100644 --- a/sevenn/train/loss.py +++ b/sevenn/train/loss.py @@ -305,7 +305,8 @@ def make_loss_info_dict_from_config(config: Dict[str, Any]): def get_loss_functions_from_config( - config: Dict[str, Any] + config: Dict[str, Any], + model_keys: Optional[List[str]] = None, ) -> List[Tuple[LossDefinition, float]]: from sevenn.train.optim import loss_dict @@ -355,4 +356,12 @@ def get_loss_functions_from_config( append_ewc_loss(loss_functions, config) + # Modal L2 regularization + # fold the 1/2 convention factor into the stored weight + if model_keys is not None: + for reg_def, reg_weight in get_regularization_from_config( + config, model_keys + ): + loss_functions.append((reg_def, reg_weight / 2.0)) + return loss_functions diff --git a/sevenn/train/trainer.py b/sevenn/train/trainer.py index 065cd399..3494b728 100644 --- a/sevenn/train/trainer.py +++ b/sevenn/train/trainer.py @@ -11,9 +11,9 @@ import sevenn._keys as KEY from sevenn.error_recorder import ErrorRecorder -from sevenn.train.loss import L2Regularization, LossDefinition +from sevenn.train.loss import LossDefinition -from .loss import get_loss_functions_from_config, get_regularization_from_config +from .loss import get_loss_functions_from_config from .optim import optim_dict, scheduler_dict @@ -38,7 +38,6 @@ def __init__( self, model: torch.nn.Module, loss_functions: List[Tuple[LossDefinition, float]], - reg_functions: Optional[List[Tuple[L2Regularization, float]]] = None, optimizer_cls=None, optimizer_args: Optional[Dict[str, Any]] = None, scheduler_cls=None, @@ -82,7 +81,6 @@ def __init__( else: self.scheduler = None self.loss_functions = loss_functions - self.reg_functions = reg_functions or [] self.grad_clip_norm_th = grad_clip_norm_th @staticmethod @@ -90,13 +88,11 @@ def from_config( model: torch.nn.Module, config: Dict[str, Any], ) -> 'Trainer': - reg_functions = get_regularization_from_config( - config, list(model._modules.keys()) - ) trainer = Trainer( model, - loss_functions=get_loss_functions_from_config(config), - reg_functions=reg_functions, + loss_functions=get_loss_functions_from_config( + config, list(model._modules.keys()) + ), optimizer_cls=optim_dict[config.get(KEY.OPTIMIZER, 'adam').lower()], optimizer_args=config.get(KEY.OPTIM_PARAM, {}), scheduler_cls=scheduler_dict[ @@ -129,8 +125,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) - reg_functions = get_regularization_from_config( + loss_functions = get_loss_functions_from_config( config, list(model._modules.keys()) ) @@ -138,7 +133,6 @@ def args_from_checkpoint(checkpoint: str) -> Tuple[Dict, Dict, Dict]: { 'model': model, 'loss_functions': loss_functions, - 'reg_functions': reg_functions, 'optimizer_cls': optimizer_cls, 'optimizer_args': config[KEY.OPTIM_PARAM], 'scheduler_cls': scheduler_cls, @@ -187,9 +181,6 @@ def run_one_epoch( indv_loss = loss_def.get_loss(output, _model) if indv_loss is not None: total_loss += (indv_loss * w) - for reg_def, w in self.reg_functions: - reg_loss = reg_def.get_loss(output, _model) - total_loss += reg_loss * w / 2 total_loss.backward() if self.grad_clip_norm_th is not None: torch.nn.utils.clip_grad_norm_( @@ -225,9 +216,6 @@ def train_one_batch( 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 - for reg_def, w in self.reg_functions: - reg_loss = reg_def.get_loss(output, _model) - total_loss += reg_loss * w / 2 total_loss.backward() From a520184ecd886a097e75047aab1b29000d90191b Mon Sep 17 00:00:00 2001 From: alphalm4 Date: Mon, 22 Jun 2026 17:34:51 +0900 Subject: [PATCH 24/32] pass only loss_functions to ErrorRecorder.from_config --- sevenn/error_recorder.py | 18 ++++-------------- sevenn/scripts/processing_by_batch.py | 2 +- sevenn/scripts/processing_epoch.py | 2 +- sevenn/train/loss.py | 2 +- 4 files changed, 7 insertions(+), 17 deletions(-) diff --git a/sevenn/error_recorder.py b/sevenn/error_recorder.py index f646459c..5caf253b 100644 --- a/sevenn/error_recorder.py +++ b/sevenn/error_recorder.py @@ -467,7 +467,6 @@ def init_total_loss_metric( def from_config( config: Dict[str, Any], loss_functions: Optional[List[Tuple[LossDefinition, float]]] = None, - reg_functions: Optional[List[Tuple[LossDefinition, float]]] = None, ) -> 'ErrorRecorder': loss_config = config.get(KEY.LOSS, 'mse') if isinstance(loss_config, dict) and loss_functions is None: @@ -482,15 +481,6 @@ def from_config( else: criteria = None - if loss_functions is not None: - all_loss_functions = ( - loss_functions + reg_functions - if isinstance(reg_functions, list) - else loss_functions - ) - else: - all_loss_functions = None - err_config = config.get(KEY.ERROR_RECORD, False) if not err_config: raise ValueError( @@ -510,14 +500,14 @@ def from_config( if err_type == 'TotalLoss': # special case err_metrics.append( ErrorRecorder.init_total_loss_metric( - config, criteria, all_loss_functions + config, criteria, loss_functions ) ) continue elif err_type == 'Modal_cos': # special case metric_cls = ModalWeightCosine metric_kwargs['loss_def'], _ = _get_loss_function_from_name( - all_loss_functions, 'L2_modal' + loss_functions, 'L2_modal' ) metric_kwargs.pop('unit', None) err_metrics.append(metric_cls(**metric_kwargs)) @@ -525,10 +515,10 @@ def from_config( metric_cls = ErrorRecorder.METRIC_DICT[metric_name] assert isinstance(metric_kwargs['name'], str) if metric_name == 'Loss': - if all_loss_functions is not None: + if loss_functions is not None: metric_cls = LossError metric_kwargs['loss_def'], _ = _get_loss_function_from_name( - all_loss_functions, metric_kwargs['name'] + loss_functions, metric_kwargs['name'] ) else: metric_cls = CustomError diff --git a/sevenn/scripts/processing_by_batch.py b/sevenn/scripts/processing_by_batch.py index 79be8772..74e65da7 100644 --- a/sevenn/scripts/processing_by_batch.py +++ b/sevenn/scripts/processing_by_batch.py @@ -59,7 +59,7 @@ def processing_by_batch( 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, trainer.reg_functions + config, trainer.loss_functions ) recorders = {k: deepcopy(recorder) for k in loaders} diff --git a/sevenn/scripts/processing_epoch.py b/sevenn/scripts/processing_epoch.py index a1212c5e..34338403 100644 --- a/sevenn/scripts/processing_epoch.py +++ b/sevenn/scripts/processing_epoch.py @@ -38,7 +38,7 @@ def processing_epoch_v2( best_metric = best_metric or config.get(KEY.BEST_METRIC, 'TotalLoss') assert isinstance(best_metric, str) recorder = error_recorder or ErrorRecorder.from_config( - config, trainer.loss_functions, trainer.reg_functions + config, trainer.loss_functions ) recorders = {k: deepcopy(recorder) for k in loaders} # reEWC: log the replayed memory set as a separate 'memoryset' column group. diff --git a/sevenn/train/loss.py b/sevenn/train/loss.py index c66553f8..a44824f6 100644 --- a/sevenn/train/loss.py +++ b/sevenn/train/loss.py @@ -261,7 +261,7 @@ def get_regularization_from_config( reg_functions: List[Tuple[LossDefinition, float]] = [] modal_param = reg_params.get('modal', {}) - if not modal_param: + if not modal_param or not config.get(KEY.USE_MODALITY, False): return reg_functions reg_weight = float(modal_param.get(KEY.REG_WEIGHT, 1e-5)) From e3e9bda1418fa11ccf60fb0fecade0f7cab2d780 Mon Sep 17 00:00:00 2001 From: alphalm4 Date: Wed, 24 Jun 2026 15:56:49 +0900 Subject: [PATCH 25/32] fix lc.csv --- sevenn/scripts/processing_by_batch.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sevenn/scripts/processing_by_batch.py b/sevenn/scripts/processing_by_batch.py index 74e65da7..f891d249 100644 --- a/sevenn/scripts/processing_by_batch.py +++ b/sevenn/scripts/processing_by_batch.py @@ -76,7 +76,7 @@ def processing_by_batch( csv_path = unique_filepath(f'{prefix}/lc.csv') if write_csv: - head = ['epoch', 'lr'] + 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: From fc99720ce4f0f3d07ac80f438c74ab00a4f0e149 Mon Sep 17 00:00:00 2001 From: alphalm4 Date: Wed, 24 Jun 2026 16:02:52 +0900 Subject: [PATCH 26/32] refactor --- sevenn/error_recorder.py | 20 ++++++++++++++++++++ sevenn/train/loss.py | 8 ++------ 2 files changed, 22 insertions(+), 6 deletions(-) diff --git a/sevenn/error_recorder.py b/sevenn/error_recorder.py index 5caf253b..8ebf4d8b 100644 --- a/sevenn/error_recorder.py +++ b/sevenn/error_recorder.py @@ -308,6 +308,18 @@ def update( 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. @@ -512,6 +524,14 @@ def from_config( 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/train/loss.py b/sevenn/train/loss.py index a44824f6..70c8dfc4 100644 --- a/sevenn/train/loss.py +++ b/sevenn/train/loss.py @@ -235,7 +235,7 @@ def get_loss( reg_params = list(module._modules['linear'].weight_views())[-1] reg_loss = torch.sum(torch.pow(reg_params, 2)) ret = ret + reg_loss - return ret + return 0.5 * ret # penalty = (1/2)||w||^2 def get_cosine( self, batch_data: Dict[str, Any], model: Optional[Callable] = None @@ -357,11 +357,7 @@ def get_loss_functions_from_config( append_ewc_loss(loss_functions, config) # Modal L2 regularization - # fold the 1/2 convention factor into the stored weight if model_keys is not None: - for reg_def, reg_weight in get_regularization_from_config( - config, model_keys - ): - loss_functions.append((reg_def, reg_weight / 2.0)) + loss_functions.extend(get_regularization_from_config(config, model_keys)) return loss_functions From 463209695346c3518b855d8c415885ca077a8af5 Mon Sep 17 00:00:00 2001 From: YutackPark Date: Fri, 26 Jun 2026 13:04:44 +0900 Subject: [PATCH 27/32] refactor --- sevenn/scripts/processing_by_batch.py | 55 +++++++++++++++++++- sevenn/scripts/processing_continue.py | 4 +- sevenn/scripts/train.py | 75 ++++----------------------- sevenn/train/aselmdb_dataset.py | 3 +- sevenn/train/loss.py | 41 +++++++-------- sevenn/train/reewc/trainer.py | 2 +- sevenn/train/trainer.py | 13 ++--- 7 files changed, 89 insertions(+), 104 deletions(-) diff --git a/sevenn/scripts/processing_by_batch.py b/sevenn/scripts/processing_by_batch.py index f891d249..92534d7d 100644 --- a/sevenn/scripts/processing_by_batch.py +++ b/sevenn/scripts/processing_by_batch.py @@ -1,7 +1,8 @@ +import math import os import time from copy import deepcopy -from typing import Optional +from typing import Any, Dict, Optional import numpy as np @@ -12,6 +13,58 @@ 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, diff --git a/sevenn/scripts/processing_continue.py b/sevenn/scripts/processing_continue.py index 01e44a6c..6bef8073 100644 --- a/sevenn/scripts/processing_continue.py +++ b/sevenn/scripts/processing_continue.py @@ -20,8 +20,7 @@ def processing_continue_v2(config: Dict[str, Any]): Skips model compatibility Returns: - For epoch training: (state_dicts, epoch) - For batch training: (state_dicts, epoch, data_progress) + (state_dicts, epoch, data_progress or None) """ log = Logger() continue_dct = config[KEY.CONTINUE] @@ -110,7 +109,6 @@ def processing_continue_v2(config: Dict[str, Any]): 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(f'Checkpoint previous epoch was: {from_epoch}') # duplicated? # noqa: E501 log.writeline('checkpoint loading success') diff --git a/sevenn/scripts/train.py b/sevenn/scripts/train.py index 2158555f..a90c5151 100644 --- a/sevenn/scripts/train.py +++ b/sevenn/scripts/train.py @@ -86,56 +86,6 @@ def loader_from_config( return DataLoader(**loader_args) -def update_config_for_batch_training(config: Dict[str, Any], train_loader) -> None: - """ - Update scheduler parameters for batch-level training. - - This converts epoch-based scheduler parameters to step-based parameters - when using batch training mode. - """ - # 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 datasets_from_py(config, script): if isinstance(script, list): assert len(script) == 1, 'Need single python script' @@ -166,7 +116,10 @@ def train_v2(config: Dict[str, Any], working_dir: str) -> None: import sevenn.train.graph_dataset as graph_dataset import sevenn.train.modal_dataset as modal_dataset - from .processing_by_batch import processing_by_batch + 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 @@ -223,39 +176,31 @@ def train_v2(config: Dict[str, Any], working_dir: str) -> None: 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['trainset']) + 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) if train_by_batch: processing_by_batch( - config, - trainer, - loaders, - data_progress, - start_epoch, - working_dir=working_dir, + config, trainer, loaders, data_progress, start_epoch, working_dir ) else: - processing_epoch_v2( - config, trainer, loaders, start_epoch, working_dir=working_dir - ) + processing_epoch_v2(config, trainer, loaders, start_epoch, working_dir) log.timer_end('total', message='Total wall time') diff --git a/sevenn/train/aselmdb_dataset.py b/sevenn/train/aselmdb_dataset.py index f875d73f..b425e324 100644 --- a/sevenn/train/aselmdb_dataset.py +++ b/sevenn/train/aselmdb_dataset.py @@ -4,7 +4,6 @@ import os import os.path as osp import time -import typing import warnings import zlib from collections import Counter @@ -16,7 +15,6 @@ import lmdb import numpy as np import orjson -import torch.distributed as dist from ase.data import chemical_symbols from ase.db.core import Database, now, ops from ase.db.row import AtomsRow @@ -42,6 +40,7 @@ class LMDBDatabase(Database): 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, diff --git a/sevenn/train/loss.py b/sevenn/train/loss.py index 70c8dfc4..941eabcd 100644 --- a/sevenn/train/loss.py +++ b/sevenn/train/loss.py @@ -225,9 +225,7 @@ def __init__( 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 - ): + 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: @@ -254,20 +252,21 @@ def get_cosine( return ret -def get_regularization_from_config( - config: Dict[str, Any], all_module_keys: List[str] -) -> List[Tuple[LossDefinition, float]]: +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, {}) - reg_functions: List[Tuple[LossDefinition, float]] = [] modal_param = reg_params.get('modal', {}) if not modal_param or not config.get(KEY.USE_MODALITY, False): - return reg_functions + return None - reg_weight = float(modal_param.get(KEY.REG_WEIGHT, 1e-5)) + if not model: + raise ValueError('modal reg is requested but model is not given.') module_keys_to_reg = [] - for module_key in all_module_keys: + for module_key in list(model._modules.keys()): for ( use_modal_module_key, modal_module_name, @@ -281,12 +280,10 @@ def get_regularization_from_config( continue module_keys_to_reg.append(module_key) - reg_functions.append(( + return ( L2Regularization('L2_modal', module_keys_to_reg, reg_modal_only=True), - reg_weight, - )) - - return reg_functions + float(modal_param.get(KEY.REG_WEIGHT, 1e-5)), + ) def make_loss_info_dict_from_config(config: Dict[str, Any]): @@ -306,9 +303,10 @@ def make_loss_info_dict_from_config(config: Dict[str, Any]): def get_loss_functions_from_config( config: Dict[str, Any], - model_keys: Optional[List[str]] = None, + 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) @@ -352,12 +350,9 @@ def get_loss_functions_from_config( loss_function = loss_function_cls(criterion=criterion, **commons) loss_functions.append((loss_function, loss_weight)) - from sevenn.train.reewc.loss import append_ewc_loss - - append_ewc_loss(loss_functions, config) - - # Modal L2 regularization - if model_keys is not None: - loss_functions.extend(get_regularization_from_config(config, model_keys)) + 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/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/trainer.py b/sevenn/train/trainer.py index 3494b728..44fe16a3 100644 --- a/sevenn/train/trainer.py +++ b/sevenn/train/trainer.py @@ -90,9 +90,7 @@ def from_config( ) -> 'Trainer': trainer = Trainer( model, - loss_functions=get_loss_functions_from_config( - config, list(model._modules.keys()) - ), + 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[ @@ -125,9 +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, list(model._modules.keys()) - ) + loss_functions = get_loss_functions_from_config(config, model) return ( { @@ -180,7 +176,7 @@ def run_one_epoch( for loss_def, w in self.loss_functions: 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() if self.grad_clip_norm_th is not None: torch.nn.utils.clip_grad_norm_( @@ -227,8 +223,7 @@ def train_one_batch( ) if norm > self.grad_clip_norm_th and self.rank == 0: print( - f'[Clipping] Grad norm {norm:.2f} into ' - f'{self.grad_clip_norm_th}', + f'[Clipping] Grad norm {norm:.2f} into {self.grad_clip_norm_th}', flush=True, ) self.optimizer.step() From 809517b3ff221a55ebd68a012630e025d61002e5 Mon Sep 17 00:00:00 2001 From: YutackPark Date: Fri, 26 Jun 2026 13:05:41 +0900 Subject: [PATCH 28/32] typo --- sevenn/train/reewc/loss.py | 17 +++++------------ 1 file changed, 5 insertions(+), 12 deletions(-) diff --git a/sevenn/train/reewc/loss.py b/sevenn/train/reewc/loss.py index e69d8933..3ff29f92 100644 --- a/sevenn/train/reewc/loss.py +++ b/sevenn/train/reewc/loss.py @@ -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) From 4c38c692802dc9ba8674281107061154e0e56b19 Mon Sep 17 00:00:00 2001 From: YutackPark Date: Fri, 26 Jun 2026 13:27:40 +0900 Subject: [PATCH 29/32] remove import in __init__ of reewc --- sevenn/scripts/train.py | 4 ++-- sevenn/train/reewc/__init__.py | 16 ---------------- sevenn/train/reewc/loss.py | 2 +- 3 files changed, 3 insertions(+), 19 deletions(-) delete mode 100644 sevenn/train/reewc/__init__.py diff --git a/sevenn/scripts/train.py b/sevenn/scripts/train.py index a90c5151..9cc54800 100644 --- a/sevenn/scripts/train.py +++ b/sevenn/scripts/train.py @@ -15,12 +15,12 @@ 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 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 3ff29f92..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 From 843433a3b45e99b9ac169880da1ae17f7591dc8b Mon Sep 17 00:00:00 2001 From: YutackPark Date: Fri, 26 Jun 2026 13:41:08 +0900 Subject: [PATCH 30/32] fix train_v2 working_dir passed positionally to processing_epoch_v2 working_dir was bound to train_loader_key, so lc.csv was written to cwd instead of working_dir. Pass it as a keyword arg. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01V97TGG8YJ8fWT6eXM75WWC --- sevenn/scripts/train.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/sevenn/scripts/train.py b/sevenn/scripts/train.py index 9cc54800..b9fdf521 100644 --- a/sevenn/scripts/train.py +++ b/sevenn/scripts/train.py @@ -200,7 +200,9 @@ def train_v2(config: Dict[str, Any], working_dir: str) -> None: config, trainer, loaders, data_progress, start_epoch, working_dir ) else: - processing_epoch_v2(config, trainer, loaders, start_epoch, working_dir) + processing_epoch_v2( + config, trainer, loaders, start_epoch, working_dir=working_dir + ) log.timer_end('total', message='Total wall time') From b2da516b065f3cb8bd435774146d7fc468c62cd4 Mon Sep 17 00:00:00 2001 From: alphalm4 Date: Fri, 10 Jul 2026 17:14:10 +0900 Subject: [PATCH 31/32] typo --- sevenn/scripts/train.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/sevenn/scripts/train.py b/sevenn/scripts/train.py index b9fdf521..8eb1dd5a 100644 --- a/sevenn/scripts/train.py +++ b/sevenn/scripts/train.py @@ -1,5 +1,4 @@ import importlib.util -import math from pathlib import Path from typing import Any, Dict, List, Optional @@ -197,7 +196,12 @@ def train_v2(config: Dict[str, Any], working_dir: str) -> None: if train_by_batch: processing_by_batch( - config, trainer, loaders, data_progress, start_epoch, working_dir + config, + trainer, + loaders, + data_progress, + start_epoch, + working_dir=working_dir, ) else: processing_epoch_v2( From 6c658bf2e1602773d4b4a9e6b51f9d91e7398fea Mon Sep 17 00:00:00 2001 From: alphalm4 Date: Fri, 10 Jul 2026 17:28:21 +0900 Subject: [PATCH 32/32] add changelog --- CHANGELOG.md | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0d7066a7..14fabea4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,8 +3,9 @@ All notable changes to this project will be documented in this file. ## [0.13.1.dev] ### Added -- L2MAE loss -- OrderedSampler, batch training +- 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`.