import math import os import pickle import random from datetime import datetime import numpy as np import torch from torch import nn from torch.nn import init from sklearn.model_selection import train_test_split import lmdb from bson import ObjectId from tqdm import tqdm from utils.lmdb_utils import get_lmdb_keys from utils.utils import get_config_id from utils.data_utils import produce_window_batches, get_collated_batch_for_key from vsm_datascience_common.cycle_database_connection.db_utils import get_cycles_collection def get_training_config(base_config: dict, model_config: dict): try: if base_config["model_class"] != model_config["model_class"]: raise Exception("Model type differ in model config and training config") except KeyError: raise Exception("Model class missing in training or model config") training_config_id = get_config_id(base_config) # hash = training_config_id[-5:] # timestamp = datetime.now().strftime("%Y%m%d-%H%M") # training_config_id = f"{timestamp}" training_dir = os.path.abspath(f"{model_config['model_dir']}/trainings/{training_config_id}") # append identifier to config training_config = base_config.copy() training_config["id"] = training_config_id training_config["training_dir"] = training_dir if not os.path.isdir(training_dir): os.makedirs(training_dir) # save training configuration with open(f"{training_dir}/training_configuration.pickle", "wb") as f: pickle.dump(training_config, f) else: # fetch config with open(f"{training_dir}/training_configuration.pickle", "rb") as f: training_config = pickle.load(f) # update paths training_config["training_dir"] = training_dir training_config["model_dir"] = model_config["model_dir"] return training_config def get_training_config_from_file(training_dir: str, base_model_dir: str, model_configuration: dict) -> dict: # fetch config with open(f"{training_dir}/training_configuration.pickle", "rb") as f: training_config = pickle.load(f) # update paths based on base directories training_config["training_dir"] = os.path.join(os.path.abspath(base_model_dir), model_configuration["id"], "trainings", training_config["id"]) return training_config def load_model_for_usage(model_configuration: dict, training_configuration: dict, sample_id: str) -> tuple[nn.Module, dict]: """ Load the model associated with the given model configuration for the usage on the current node, given its training configuration If the model does not exist, create it Args: model_configuration: configuration of the model training_configuration: training configuration sample_id: id of a sample data item Returns: model: model training_config: training configuration with updated variables such as batch size """ # get computation rank if torch.distributed.is_initialized(): local_rank = torch.distributed.get_rank() torch.cuda.set_device(local_rank) device = torch.device(f"cuda:{local_rank}") else: local_rank = 0 device = torch.device("cuda" if torch.cuda.is_available() else "cpu") dataset_dir = model_configuration["feature_config"]["dataset_dir"] lmdb_env = lmdb.open(dataset_dir, readonly=True) try: load_fn = model_configuration["model_load_fn"] model = load_fn(model_configuration=model_configuration, training_configuration=training_configuration, sample_key=sample_id, device=device, lmdb_env=lmdb_env) except FileNotFoundError: model_creation_fn = model_configuration["model_creation_fn"] model = model_creation_fn( model_configuration=model_configuration, sample_key=sample_id, lmdb_env=lmdb_env ) # estimate batch size sample_batch = get_collated_batch_for_key(sample_id, model_configuration, lmdb_env=lmdb_env) if sample_batch is None or len(sample_batch) == 0: raise ValueError( f"Rank {local_rank}: Could not get sample input for key {sample_id} in dataset {dataset_dir}. " f"Please check if the key exists in the LMDB dataset.") sample_input = sample_batch[0][0] # get the first item in the batch input_shape = sample_input.shape # send model to proper device model.to(device) batch_size = find_max_batch_size(model, input_shape, device=device) # update max lr in training configuration, use square root scaling law based on initial batch size training_configuration["learning_parameters"]["learning_rate"] = training_configuration["learning_parameters"][ "max_lr"] * math.sqrt( batch_size / training_configuration["batch_size"]) # update batch size in training configuration training_configuration["batch_size"] = batch_size # close lmdb env lmdb_env.close() return model, training_configuration def weight_init(m): """ Usage: model = Model() model.apply(weight_init) """ if isinstance(m, nn.Conv1d): init.normal_(m.weight.data) if m.bias is not None: init.normal_(m.bias.data) elif isinstance(m, nn.Conv2d): init.xavier_normal_(m.weight.data) if m.bias is not None: init.normal_(m.bias.data) elif isinstance(m, nn.Conv3d): init.xavier_normal_(m.weight.data) if m.bias is not None: init.normal_(m.bias.data) elif isinstance(m, nn.ConvTranspose1d): init.normal_(m.weight.data) if m.bias is not None: init.normal_(m.bias.data) elif isinstance(m, nn.ConvTranspose2d): init.xavier_normal_(m.weight.data) if m.bias is not None: init.normal_(m.bias.data) elif isinstance(m, nn.ConvTranspose3d): init.xavier_normal_(m.weight.data) if m.bias is not None: init.normal_(m.bias.data) elif isinstance(m, nn.BatchNorm1d): init.normal_(m.weight.data, mean=1, std=0.02) init.constant_(m.bias.data, 0) elif isinstance(m, nn.BatchNorm2d): init.normal_(m.weight.data, mean=1, std=0.02) init.constant_(m.bias.data, 0) elif isinstance(m, nn.BatchNorm3d): init.normal_(m.weight.data, mean=1, std=0.02) init.constant_(m.bias.data, 0) elif isinstance(m, nn.Linear): init.xavier_normal_(m.weight.data) if m.bias is not None: init.normal_(m.bias.data) elif isinstance(m, nn.LSTM): for param in m.parameters(): if len(param.shape) >= 2: init.orthogonal_(param.data) else: init.normal_(param.data) elif isinstance(m, nn.LSTMCell): for param in m.parameters(): if len(param.shape) >= 2: init.orthogonal_(param.data) else: init.normal_(param.data) elif isinstance(m, nn.GRU): for param in m.parameters(): if len(param.shape) >= 2: init.orthogonal_(param.data) else: init.normal_(param.data) for names in m._all_weights: for name in filter(lambda n: "bias" in n, names): bias = getattr(m, name) n = bias.size(0) bias.data[:n // 3].fill_(-1.) elif isinstance(m, nn.GRUCell): for param in m.parameters(): if len(param.shape) >= 2: init.orthogonal_(param.data) else: init.normal_(param.data) def collate(batch_items: list) -> dict: batch = dict() for key in batch_items[0].keys(): if key in ["combination_id", "time_index"]: continue else: if batch_items[0][key] is None: batch[key] = None else: batch[key] = np.stack([item[key] for item in batch_items]) for key in batch.keys(): if batch[key] is not None: batch[key] = torch.tensor(batch[key], dtype=torch.float32) return batch def get_splits_by_user(input_keys: list, train_size: float, val_size: float, test_size: float): if train_size + val_size + test_size != 1: raise ValueError("Train, val and test sizes must sum to 1") if len(input_keys) == 0: raise ValueError("Input keys list is empty") items_by_use = dict() for input_key in tqdm(input_keys): try: user_id = get_cycles_collection().find_one({"_id": ObjectId(input_key)})["user_id"] if user_id not in items_by_use: items_by_use[user_id] = [] items_by_use[user_id].append(input_key) except Exception: print(f"Error getting user id for key {input_key}") continue user_ids = list(items_by_use.keys()) random.shuffle(user_ids) train_users, temp_users = train_test_split(user_ids, train_size=train_size, test_size=test_size + val_size) # compute relative test size, as it must be relative to the remaining users relative_test_size = test_size / (1 - train_size) val_users, test_users = train_test_split(temp_users, test_size=relative_test_size) train_keys = [] val_keys = [] test_keys = [] for user_id in train_users: train_keys.extend(items_by_use[user_id]) for user_id in val_users: val_keys.extend(items_by_use[user_id]) for user_id in test_users: test_keys.extend(items_by_use[user_id]) return train_keys, val_keys, test_keys def save_splits(train_keys: list, val_keys: list, test_keys: list, base_dir: str): if not os.path.exists(base_dir): os.makedirs(base_dir) with open(f"{base_dir}/train_keys.pickle", "wb") as f: pickle.dump(train_keys, f) with open(f"{base_dir}/val_keys.pickle", "wb") as f: pickle.dump(val_keys, f) with open(f"{base_dir}/test_keys.pickle", "wb") as f: pickle.dump(test_keys, f) def load_splits(base_dir: str): with open(f"{base_dir}/train_keys.pickle", "rb") as f: train_keys = pickle.load(f) with open(f"{base_dir}/val_keys.pickle", "rb") as f: val_keys = pickle.load(f) with open(f"{base_dir}/test_keys.pickle", "rb") as f: test_keys = pickle.load(f) return train_keys, val_keys, test_keys def get_data_ids(model_configuration: dict, training_configuration: dict, env_path: str, limit: int = None) -> tuple: env = lmdb.open(f"{env_path}", readonly=True) if os.path.exists(f"{model_configuration['feature_config']['dataset_dir']}/train_keys.pickle"): train_ids, val_ids, test_ids = load_splits(model_configuration["feature_config"]["dataset_dir"]) else: lmdb_keys = get_lmdb_keys(env, limit) # train_ids, val_ids, test_ids = get_splits_by_user(lmdb_keys, # training_configuration["train_size"], # training_configuration["val_size"], # training_configuration["test_size"]) train_ids, temp_ids = train_test_split(lmdb_keys, train_size=training_configuration["train_size"], test_size=training_configuration["val_size"] + training_configuration[ "test_size"]) # compute relative test size, as it must be relative to the remaining users relative_test_size = training_configuration["test_size"] / (1 - training_configuration["train_size"]) val_ids, test_ids = train_test_split(temp_ids, test_size=relative_test_size) # save splits to file save_splits(train_ids, val_ids, test_ids, model_configuration["feature_config"]["dataset_dir"]) if limit is not None: train_size = training_configuration["train_size"] val_size = training_configuration["val_size"] test_size = training_configuration["test_size"] train_lim = int(limit * train_size) val_lim = int(limit * val_size) test_lim = int(limit * test_size) train_ids = train_ids[:train_lim] val_ids = val_ids[:val_lim] test_ids = test_ids[:test_lim] return train_ids, val_ids, test_ids def find_max_batch_size(model, input_shape, device, dtype=torch.float32, max_batch=2048): """ Find the maximum batch size that can be processed by the model without running out of memory. Model must already be moved to the appropriate device (e.g., GPU). Args: model: model to test input_shape: shape of the input tensor (excluding batch size) device: device to work on dtype: data type of the input tensor (default: torch.float32) max_batch: maximum batch size to test (default: 2048) Returns: last_successful: the largest batch size that did not cause an out of memory error """ batch_size = 1 last_successful = 1 while last_successful < max_batch: try: inputs = torch.randn((batch_size, *input_shape), device=device, dtype=dtype) outputs = model(inputs) loss = outputs.sum() loss.backward() # simulate training torch.cuda.empty_cache() last_successful = batch_size batch_size *= 2 except RuntimeError as e: if 'out of memory' in str(e): torch.cuda.empty_cache() break else: raise e return last_successful