Source code for dicee.executer

import json
import logging
import time
import warnings
from types import SimpleNamespace
import os
import datetime
from pytorch_lightning import seed_everything
from .knowledge_graph import KG
from .evaluator import Evaluator
from .static_preprocess_funcs import preprocesses_input_args
from .trainer import DICE_Trainer
from .static_funcs import timeit, read_or_load_kg, load_json, store, create_experiment_folder
import numpy as np
import torch
import torch.distributed as dist
from pytorch_lightning.utilities.rank_zero import rank_zero_only

logging.getLogger('pytorch_lightning').setLevel(0)
warnings.filterwarnings(action="ignore", category=DeprecationWarning)
os.environ["TORCH_DISTRIBUTED_DEBUG"] = "INFO"

[docs] class Execute: """ A class for Training, Retraining and Evaluation a model. (1) Loading & Preprocessing & Serializing input data. (2) Training & Validation & Testing (3) Storing all necessary info """ def __init__(self, args, continuous_training=False): # check if we need distributed training self.distributed = getattr(args, "trainer", None) == "torchDDP" # initialize distributed training if required if self.distributed: if not dist.is_initialized(): dist.init_process_group(backend="nccl", init_method="env://") self.rank = dist.get_rank() self.world_size = dist.get_world_size() self.local_rank = int(os.environ["LOCAL_RANK"]) torch.cuda.set_device(self.local_rank) print(f"[Rank {self.rank}] mapped to GPU {self.local_rank}", flush=True) else: self.rank, self.world_size, self.local_rank = 0, 1, 0 # (1) Process arguments and sanity checking. self.args = preprocesses_input_args(args) # (2) Ensure reproducibility. seed_everything(args.random_seed, workers=True) # (3) Set the continual training flag self.is_continual_training = continuous_training # (4) Create an experiment folder or use the previous one if self.rank == 0: self.setup_executor() # (5) A variable is initialized for pytorch lightning trainer or DICE_Trainer() self.trainer = None self.trained_model = None # (6) A variable is initialized for storing input data. self.knowledge_graph = None # (7) Store few data in memory for numerical results, e.g. runtime, H@1 etc. self.report = dict() # (8) Create an object to carry out link prediction evaluations, e.g. Evaluator(self) self.evaluator = None # (9) Execution start time self.start_time = None
[docs] def is_rank_zero(self) -> bool: return self.rank == 0
[docs] def cleanup(self): if self.distributed and dist.is_initialized(): dist.destroy_process_group()
[docs] @rank_zero_only def setup_executor(self) -> None: # Only set up storage if not in continual training mode if self.is_continual_training is False: # If a specific directory for this run is provided if self.args.path_to_store_single_run is not None: # Check if we should reuse an existing directory or start fresh reuse_existing = getattr(self.args, "reuse_existing_run_dir", False) if os.path.exists(self.args.path_to_store_single_run): if not reuse_existing: # Delete the existing directory to start with a clean slate print(f"Deleting the existing directory of {self.args.path_to_store_single_run}") import shutil shutil.rmtree(self.args.path_to_store_single_run) # Create a new directory for this run os.makedirs(self.args.path_to_store_single_run, exist_ok=False) else: # Reuse the existing directory if allowed print(f"Reusing the existing directory of {self.args.path_to_store_single_run}") else: # Create the directory if it does not exist os.makedirs(self.args.path_to_store_single_run, exist_ok=False) # Set the full storage path to the specified directory self.args.full_storage_path = self.args.path_to_store_single_run else: # If no directory is specified, create a new experiment folder self.args.full_storage_path = create_experiment_folder(folder_name=self.args.storage_path) # Update the path for this run to the newly created folder self.args.path_to_store_single_run = self.args.full_storage_path # Save the current configuration to a JSON file in the storage directory with open(self.args.full_storage_path + '/configuration.json', 'w') as file_descriptor: temp = vars(self.args) json.dump(temp, file_descriptor, indent=3)
[docs] def create_and_store_kg(self): if not self.is_rank_zero(): return memmap_path = os.path.join(self.args.path_to_store_single_run, "memory_map_train_set.npy") details_path = os.path.join(self.args.path_to_store_single_run, "memory_map_details.json") if os.path.exists(memmap_path) and os.path.exists(details_path): print("KG memmap already exists, skipping.") return print("Creating knowledge graph...") self.knowledge_graph = read_or_load_kg(self.args, cls=KG) self.args.num_entities = self.knowledge_graph.num_entities self.args.num_relations = self.knowledge_graph.num_relations self.args.num_tokens = self.knowledge_graph.num_tokens self.args.max_length_subword_tokens = self.knowledge_graph.max_length_subword_tokens self.args.ordered_bpe_entities = self.knowledge_graph.ordered_bpe_entities self.report['num_train_triples'] = len(self.knowledge_graph.train_set) self.report['num_entities'] = self.knowledge_graph.num_entities self.report['num_relations'] = self.knowledge_graph.num_relations self.report['max_length_subword_tokens'] = self.knowledge_graph.max_length_subword_tokens if self.knowledge_graph.max_length_subword_tokens else None self.report['runtime_kg_loading'] = time.time() - self.start_time data={"shape":tuple(self.knowledge_graph.train_set.shape), "dtype":self.knowledge_graph.train_set.dtype.str, "num_entities":self.knowledge_graph.num_entities, "num_relations":self.knowledge_graph.num_relations} with open(self.args.full_storage_path + '/memory_map_details.json', 'w') as file_descriptor: json.dump(data, file_descriptor, indent=4) memmap_kg = np.memmap(memmap_path, dtype=self.knowledge_graph.train_set.dtype, mode='w+', shape=self.knowledge_graph.train_set.shape) memmap_kg[:] = self.knowledge_graph.train_set[:] memmap_kg.flush() del memmap_kg
[docs] def load_from_memmap(self): with open(self.args.path_to_store_single_run+'/memory_map_details.json', 'r') as file_descriptor: memory_map_details = json.load(file_descriptor) self.knowledge_graph = np.memmap(self.args.path_to_store_single_run + '/memory_map_train_set.npy', mode='r', dtype=memory_map_details["dtype"], shape=tuple(memory_map_details["shape"])) self.args.num_entities = memory_map_details["num_entities"] self.args.num_relations = memory_map_details["num_relations"] self.args.num_tokens = None self.args.max_length_subword_tokens = None self.args.ordered_bpe_entities = None
[docs] @timeit def save_trained_model(self) -> None: """ Save a knowledge graph embedding model (1) Send model to eval mode and cpu. (2) Store the memory footprint of the model. (3) Save the model into disk. (4) Update the stats of KG again ? Parameter ---------- Return ---------- None """ print('*** Save Trained Model ***') self.trained_model.eval() self.trained_model.to('cpu') # Save the epoch loss # (2) Store NumParam and EstimatedSizeMB self.report.update(self.trained_model.mem_of_model()) # (3) Store/Serialize Model for further use. if self.is_continual_training is False: store(trained_model=self.trained_model, model_name='model', full_storage_path=self.args.full_storage_path, save_embeddings_as_csv=self.args.save_embeddings_as_csv) else: store(trained_model=self.trained_model, model_name='model', # + str(datetime.datetime.now()), full_storage_path=self.args.full_storage_path, save_embeddings_as_csv=self.args.save_embeddings_as_csv) self.report['path_experiment_folder'] = self.args.full_storage_path self.report['num_entities'] = self.args.num_entities self.report['num_relations'] = self.args.num_relations
[docs] @rank_zero_only def end(self, form_of_labelling: str) -> dict: """ End training (1) Store trained model. (2) Report runtimes. (3) Eval model if required. Parameter --------- Returns ------- A dict containing information about the training and/or evaluation """ # (1) Save the model self.save_trained_model() # (2) Report self.write_report() # (3) Eval model and return eval results. if self.args.eval_model is None: self.write_report() return {**self.report} else: self.evaluator.eval(dataset=self.knowledge_graph, trained_model=self.trained_model, form_of_labelling=form_of_labelling) self.write_report() return {**self.report, **self.evaluator.report}
[docs] def write_report(self) -> None: """ Report training related information in a report.json file """ # @TODO: Move to static funcs # Report total runtime. self.report['Runtime'] = time.time() - self.start_time print(f"Total Runtime: {self.report['Runtime']:.3f} seconds") with open(self.args.full_storage_path + '/report.json', 'w') as file_descriptor: json.dump(self.report, file_descriptor, indent=4)
[docs] def start(self) -> dict: """ Start training # (1) Loading the Data # (2) Create an evaluator object. # (3) Create a trainer object. # (4) Start the training Parameter --------- Returns ------- A dict containing information about the training and/or evaluation """ self.start_time = time.time() print(f"Start time:{datetime.datetime.now()}") # (1) Create knowledge graph self.create_and_store_kg() # (2) Synchronize processes if distributed training is used if self.distributed and dist.is_initialized(): dist.barrier() # (3) Reload the memory-map of index knowledge graph stored as a numpy ndarray if self.knowledge_graph is None: self.load_from_memmap() # (4) Create an evaluator object. self.evaluator = Evaluator(args=self.args) # (5) Create a trainer object. if not getattr(self.args, "full_storage_path", None): self.args.full_storage_path = self.args.path_to_store_single_run self.trainer = DICE_Trainer(args=self.args, is_continual_training=self.is_continual_training, storage_path=self.args.full_storage_path, evaluator=self.evaluator) # (6) Start the training self.trained_model, form_of_labelling = self.trainer.start(knowledge_graph=self.knowledge_graph) return self.end(form_of_labelling)
[docs] class ContinuousExecute(Execute): """ A subclass of Execute Class for retraining (1) Loading & Preprocessing & Serializing input data. (2) Training & Validation & Testing (3) Storing all necessary info During the continual learning we can only modify *** num_epochs *** parameter. Trained model stored in the same folder as the seed model for the training. Trained model is noted with the current time. """ def __init__(self, args): # (1) Current input configuration. assert os.path.exists(args.continual_learning), f"Path doesn't exist {args.continual_learning}" assert os.path.isfile(args.continual_learning + '/configuration.json') # (2) Load previous input configuration. previous_args = load_json(args.continual_learning + '/configuration.json') args=vars(args) # previous_args["num_epochs"]=args["num_epochs"] previous_args["continual_learning"]=args["continual_learning"] print("Updated configuration:",previous_args) try: report = load_json(args['continual_learning'] + '/report.json') previous_args['num_entities'] = report['num_entities'] previous_args['num_relations'] = report['num_relations'] except AssertionError: print("Couldn't find report.json.") previous_args = SimpleNamespace(**previous_args) print('ContinuousExecute starting...') print(previous_args) super().__init__(previous_args, continuous_training=True)
[docs] def continual_start(self) -> dict: """ Start Continual Training (1) Initialize training. (2) Start continual training. (3) Save trained model. Parameter --------- Returns ------- A dict containing information about the training and/or evaluation """ # (1) self.trainer = DICE_Trainer(args=self.args, is_continual_training=True, storage_path=self.args.continual_learning) # (2) assert os.path.exists(f"{self.args.continual_learning}/memory_map_train_set.npy") # (1) Reload the memory-map of index knowledge graph stored as a numpy ndarray. with open(f"{self.args.continual_learning}/memory_map_details.json", 'r') as file_descriptor: memory_map_details = json.load(file_descriptor) knowledge_graph = np.memmap(f"{self.args.continual_learning}/memory_map_train_set.npy", mode='r', dtype=memory_map_details["dtype"], shape=tuple(memory_map_details["shape"])) self.args.num_entities = memory_map_details["num_entities"] self.args.num_relations = memory_map_details["num_relations"] self.args.num_tokens = None self.args.max_length_subword_tokens = None self.args.ordered_bpe_entities = None self.trained_model, form_of_labelling = self.trainer.continual_start(knowledge_graph) # (5) Store trained model. self.save_trained_model() # (6) Eval model. if self.args.eval_model is None: return self.report else: self.evaluator = Evaluator(args=self.args, is_continual_training=True) self.evaluator.dummy_eval(self.trained_model, form_of_labelling) return {**self.report, **self.evaluator.report}