train.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up logger\r\n",
    "import os\r\n",
    "import logging\r\n",
    "\r\n",
    "os.environ[\"TF_CPP_MIN_LOG_LEVEL\"] = \"1\"\r\n",
    "os.environ[\"AUTOGRAPH_VERBOSITY\"] = \"1\"\r\n",
    "\r\n",
    "logger = logging.getLogger()\r\n",
    "logger.handlers = []\r\n",
    "ch = logging.StreamHandler()\r\n",
    "formatter = logging.Formatter(\r\n",
    "    fmt=\"%(asctime)s (%(levelname)s): %(message)s\", datefmt=\"%Y-%m-%d %H:%M:%S\"\r\n",
    ")\r\n",
    "ch.setFormatter(formatter)\r\n",
    "logger.addHandler(ch)\r\n",
    "logger.setLevel(\"INFO\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\r\n",
    "import yaml\r\n",
    "import string\r\n",
    "import ast\r\n",
    "import random\r\n",
    "import time\r\n",
    "from datetime import datetime\r\n",
    "\r\n",
    "from gemnet.model.gemnet import GemNet\r\n",
    "from gemnet.training.trainer import Trainer\r\n",
    "from gemnet.training.metrics import Metrics, BestMetrics\r\n",
    "from gemnet.training.data_container import DataContainer\r\n",
    "from gemnet.training.data_provider import DataProvider\r\n",
    "\r\n",
    "import torch\r\n",
    "from torch.utils.tensorboard import SummaryWriter"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load config file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('config.yaml', 'r') as c:\r\n",
    "    config = yaml.safe_load(c)\r\n",
    "    \r\n",
    "# For strings that yaml doesn't parse (e.g. None)\r\n",
    "for key, val in config.items():\r\n",
    "    if type(val) is str:\r\n",
    "        try:\r\n",
    "            config[key] = ast.literal_eval(val)\r\n",
    "        except (ValueError, SyntaxError):\r\n",
    "            pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_spherical = config[\"num_spherical\"]\r\n",
    "num_radial = config[\"num_radial\"]\r\n",
    "num_blocks = config[\"num_blocks\"]\r\n",
    "emb_size_atom = config[\"emb_size_atom\"]\r\n",
    "emb_size_edge = config[\"emb_size_edge\"]\r\n",
    "emb_size_trip = config[\"emb_size_trip\"]\r\n",
    "emb_size_quad = config[\"emb_size_quad\"]\r\n",
    "emb_size_rbf = config[\"emb_size_rbf\"]\r\n",
    "emb_size_cbf = config[\"emb_size_cbf\"]\r\n",
    "emb_size_sbf = config[\"emb_size_sbf\"]\r\n",
    "num_before_skip = config[\"num_before_skip\"]\r\n",
    "num_after_skip = config[\"num_after_skip\"]\r\n",
    "num_concat = config[\"num_concat\"]\r\n",
    "num_atom = config[\"num_atom\"]\r\n",
    "emb_size_bil_quad = config[\"emb_size_bil_quad\"]\r\n",
    "emb_size_bil_trip = config[\"emb_size_bil_trip\"]\r\n",
    "triplets_only = config[\"triplets_only\"]\r\n",
    "forces_coupled = config[\"forces_coupled\"]\r\n",
    "direct_forces = config[\"direct_forces\"]\r\n",
    "mve = config[\"mve\"]\r\n",
    "cutoff = config[\"cutoff\"]\r\n",
    "int_cutoff = config[\"int_cutoff\"]\r\n",
    "envelope_exponent = config[\"envelope_exponent\"]\r\n",
    "extensive = config[\"extensive\"]\r\n",
    "output_init = config[\"output_init\"]\r\n",
    "scale_file = config[\"scale_file\"]\r\n",
    "data_seed = config[\"data_seed\"]\r\n",
    "dataset = config[\"dataset\"]\r\n",
    "val_dataset = config[\"val_dataset\"]\r\n",
    "num_train = config[\"num_train\"]\r\n",
    "num_val = config[\"num_val\"]\r\n",
    "logdir = config[\"logdir\"]\r\n",
    "loss = config[\"loss\"]\r\n",
    "tfseed = config[\"tfseed\"]\r\n",
    "num_steps = config[\"num_steps\"]\r\n",
    "rho_force = config[\"rho_force\"]\r\n",
    "ema_decay = config[\"ema_decay\"]\r\n",
    "weight_decay = config[\"weight_decay\"]\r\n",
    "grad_clip_max = config[\"grad_clip_max\"]\r\n",
    "agc = config[\"agc\"]\r\n",
    "decay_patience = config[\"decay_patience\"]\r\n",
    "decay_factor = config[\"decay_factor\"]\r\n",
    "decay_cooldown = config[\"decay_cooldown\"]\r\n",
    "batch_size = config[\"batch_size\"]\r\n",
    "evaluation_interval = config[\"evaluation_interval\"]\r\n",
    "patience = config[\"patience\"]\r\n",
    "save_interval = config[\"save_interval\"]\r\n",
    "learning_rate = config[\"learning_rate\"]\r\n",
    "warmup_steps = config[\"warmup_steps\"]\r\n",
    "decay_steps = config[\"decay_steps\"]\r\n",
    "decay_rate = config[\"decay_rate\"]\r\n",
    "staircase = config[\"staircase\"]\r\n",
    "restart = config[\"restart\"]\r\n",
    "comment = config[\"comment\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Set paths and create directories"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(tfseed)\r\n",
    "\r\n",
    "logging.info(\"Start training\")\r\n",
    "num_gpus = torch.cuda.device_count()\r\n",
    "cuda_available = torch.cuda.is_available()\r\n",
    "logging.info(f\"Available GPUs: {num_gpus}\")\r\n",
    "logging.info(f\"CUDA Available: {cuda_available}\")\r\n",
    "if num_gpus == 0:\r\n",
    "    logging.warning(\"No GPUs were found. Training is run on CPU!\")\r\n",
    "if not cuda_available:\r\n",
    "    logging.warning(\"CUDA unavailable. Training is run on CPU!\")\r\n",
    "\r\n",
    "# Used for creating a \"unique\" id for a run (almost impossible to generate the same twice)\r\n",
    "def id_generator(\r\n",
    "    size=6, chars=string.ascii_uppercase + string.ascii_lowercase + string.digits\r\n",
    "):\r\n",
    "    return \"\".join(random.SystemRandom().choice(chars) for _ in range(size))\r\n",
    "\r\n",
    "# A unique directory name is created for this run based on the input\r\n",
    "if (restart is None) or (restart == \"None\"):\r\n",
    "    directory = (\r\n",
    "        logdir\r\n",
    "        + \"/\"\r\n",
    "        + datetime.now().strftime(\"%Y%m%d_%H%M%S\")\r\n",
    "        + \"_\"\r\n",
    "        + id_generator()\r\n",
    "        + \"_\"\r\n",
    "        + os.path.basename(dataset)\r\n",
    "        + \"_\"\r\n",
    "        + str(comment)\r\n",
    "    )\r\n",
    "else:\r\n",
    "    directory = restart\r\n",
    "\r\n",
    "logging.info(f\"Directory: {directory}\")\r\n",
    "logging.info(\"Create directories\")\r\n",
    "\r\n",
    "if not os.path.exists(directory):\r\n",
    "    os.makedirs(directory, exist_ok=True)\r\n",
    "\r\n",
    "best_dir = os.path.join(directory, \"best\")\r\n",
    "if not os.path.exists(best_dir):\r\n",
    "    os.makedirs(best_dir)\r\n",
    "log_dir = os.path.join(directory, \"logs\")\r\n",
    "if not os.path.exists(log_dir):\r\n",
    "    os.makedirs(log_dir)\r\n",
    "\r\n",
    "extension = \".pth\"\r\n",
    "log_path_model = f\"{log_dir}/model{extension}\"\r\n",
    "log_path_training = f\"{log_dir}/training{extension}\"\r\n",
    "best_path_model = f\"{best_dir}/model{extension}\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Initialize model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "logging.info(\"Initialize model\")\r\n",
    "model = GemNet(\r\n",
    "    num_spherical=num_spherical,\r\n",
    "    num_radial=num_radial,\r\n",
    "    num_blocks=num_blocks,\r\n",
    "    emb_size_atom=emb_size_atom,\r\n",
    "    emb_size_edge=emb_size_edge,\r\n",
    "    emb_size_trip=emb_size_trip,\r\n",
    "    emb_size_quad=emb_size_quad,\r\n",
    "    emb_size_rbf=emb_size_rbf,\r\n",
    "    emb_size_cbf=emb_size_cbf,\r\n",
    "    emb_size_sbf=emb_size_sbf,\r\n",
    "    num_before_skip=num_before_skip,\r\n",
    "    num_after_skip=num_after_skip,\r\n",
    "    num_concat=num_concat,\r\n",
    "    num_atom=num_atom,\r\n",
    "    emb_size_bil_quad=emb_size_bil_quad,\r\n",
    "    emb_size_bil_trip=emb_size_bil_trip,\r\n",
    "    num_targets=2 if mve else 1,\r\n",
    "    triplets_only=triplets_only,\r\n",
    "    direct_forces=direct_forces,\r\n",
    "    forces_coupled=forces_coupled,\r\n",
    "    cutoff=cutoff,\r\n",
    "    int_cutoff=int_cutoff,\r\n",
    "    envelope_exponent=envelope_exponent,\r\n",
    "    activation=\"swish\",\r\n",
    "    extensive=extensive,\r\n",
    "    output_init=output_init,\r\n",
    "    scale_file=scale_file,\r\n",
    ")\r\n",
    "# push to GPU if available\r\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\r\n",
    "model.to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train = {}\r\n",
    "validation = {}\r\n",
    "\r\n",
    "logging.info(\"Load dataset\")\r\n",
    "data_container = DataContainer(\r\n",
    "    dataset, cutoff=cutoff, int_cutoff=int_cutoff, triplets_only=triplets_only\r\n",
    ")\r\n",
    "\r\n",
    "if val_dataset is not None:\r\n",
    "    # Initialize DataProvider\r\n",
    "    if num_train == 0:\r\n",
    "        num_train = len(data_container)\r\n",
    "    logging.info(f\"Training data size: {num_train}\")\r\n",
    "    data_provider = DataProvider(\r\n",
    "        data_container,\r\n",
    "        num_train,\r\n",
    "        0,\r\n",
    "        batch_size,\r\n",
    "        seed=data_seed,\r\n",
    "        shuffle=True,\r\n",
    "        random_split=True,\r\n",
    "    )\r\n",
    "\r\n",
    "    # Initialize validation datasets\r\n",
    "    val_data_container = DataContainer(\r\n",
    "        val_dataset,\r\n",
    "        cutoff=cutoff,\r\n",
    "        int_cutoff=int_cutoff,\r\n",
    "        triplets_only=triplets_only,\r\n",
    "    )\r\n",
    "    if num_val == 0:\r\n",
    "        num_val = len(val_data_container)\r\n",
    "    logging.info(f\"Validation data size: {num_val}\")\r\n",
    "    val_data_provider = DataProvider(\r\n",
    "        val_data_container,\r\n",
    "        0,\r\n",
    "        num_val,\r\n",
    "        batch_size,\r\n",
    "        seed=data_seed,\r\n",
    "        shuffle=True,\r\n",
    "        random_split=True,\r\n",
    "    )\r\n",
    "else:\r\n",
    "    # Initialize DataProvider (splits dataset into 3 sets based on data_seed and provides tf.datasets)\r\n",
    "    logging.info(f\"Training data size: {num_train}\")\r\n",
    "    logging.info(f\"Validation data size: {num_val}\")\r\n",
    "    assert num_train > 0\r\n",
    "    assert num_val > 0\r\n",
    "    data_provider = DataProvider(\r\n",
    "        data_container,\r\n",
    "        num_train,\r\n",
    "        num_val,\r\n",
    "        batch_size,\r\n",
    "        seed=data_seed,\r\n",
    "        shuffle=True,\r\n",
    "        random_split=True,\r\n",
    "    )\r\n",
    "    val_data_provider = data_provider\r\n",
    "\r\n",
    "# Initialize datasets\r\n",
    "train[\"dataset_iter\"] = data_provider.get_dataset(\"train\")\r\n",
    "validation[\"dataset_iter\"] = val_data_provider.get_dataset(\"val\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Prepare training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "logging.info(\"Prepare training\")\r\n",
    "# Initialize trainer\r\n",
    "trainer = Trainer(\r\n",
    "    model,\r\n",
    "    learning_rate=learning_rate,\r\n",
    "    decay_steps=decay_steps,\r\n",
    "    decay_rate=decay_rate,\r\n",
    "    warmup_steps=warmup_steps,\r\n",
    "    weight_decay=weight_decay,\r\n",
    "    ema_decay=ema_decay,\r\n",
    "    decay_patience=decay_patience,\r\n",
    "    decay_factor=decay_factor,\r\n",
    "    decay_cooldown=decay_cooldown,\r\n",
    "    grad_clip_max=grad_clip_max,\r\n",
    "    rho_force=rho_force,\r\n",
    "    mve=mve,\r\n",
    "    loss=loss,\r\n",
    "    staircase=staircase,\r\n",
    "    agc=agc,\r\n",
    ")\r\n",
    "\r\n",
    "# Initialize metrics\r\n",
    "train[\"metrics\"] = Metrics(\"train\", trainer.tracked_metrics)\r\n",
    "validation[\"metrics\"] = Metrics(\"val\", trainer.tracked_metrics)\r\n",
    "\r\n",
    "# Save/load best recorded loss (only the best model is saved)\r\n",
    "metrics_best = BestMetrics(best_dir, validation[\"metrics\"])\r\n",
    "\r\n",
    "# Set up checkpointing\r\n",
    "# Restore latest checkpoint\r\n",
    "if os.path.exists(log_path_model):\r\n",
    "    logging.info(\"Restoring model and trainer\")\r\n",
    "    model_checkpoint = torch.load(log_path_model)\r\n",
    "    model.load_state_dict(model_checkpoint[\"model\"])\r\n",
    "\r\n",
    "    train_checkpoint = torch.load(log_path_training)\r\n",
    "    trainer.load_state_dict(train_checkpoint[\"trainer\"])\r\n",
    "    # restore the best saved results\r\n",
    "    metrics_best.restore()\r\n",
    "    logging.info(f\"Restored best metrics: {metrics_best.loss}\")\r\n",
    "    step_init = int(train_checkpoint[\"step\"])\r\n",
    "else:\r\n",
    "    logging.info(\"Freshly initialize model\")\r\n",
    "    metrics_best.inititalize()\r\n",
    "    step_init = 0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "summary_writer = SummaryWriter(log_dir)\r\n",
    "steps_per_epoch = int(np.ceil(num_train / batch_size))\r\n",
    "\r\n",
    "for step in range(step_init + 1, num_steps + 1):\r\n",
    "\r\n",
    "    # keep track of the learning rate\r\n",
    "    if step % 10 == 0:\r\n",
    "        lr = trainer.schedulers[0].get_last_lr()[0]\r\n",
    "        summary_writer.add_scalar(\"lr\", lr, global_step=step)\r\n",
    "\r\n",
    "    # Perform training step\r\n",
    "    trainer.train_on_batch(train[\"dataset_iter\"], train[\"metrics\"])\r\n",
    "\r\n",
    "    # Save progress\r\n",
    "    if step % save_interval == 0:\r\n",
    "        torch.save({\"model\": model.state_dict()}, log_path_model)\r\n",
    "        torch.save(\r\n",
    "            {\"trainer\": trainer.state_dict(), \"step\": step}, log_path_training\r\n",
    "        )\r\n",
    "\r\n",
    "    # Check performance on the validation set\r\n",
    "    if step % evaluation_interval == 0:\r\n",
    "\r\n",
    "        # Save backup variables and load averaged variables\r\n",
    "        trainer.save_variable_backups()\r\n",
    "        trainer.load_averaged_variables()\r\n",
    "\r\n",
    "        # Compute averages\r\n",
    "        for i in range(int(np.ceil(num_val / batch_size))):\r\n",
    "            trainer.test_on_batch(validation[\"dataset_iter\"], validation[\"metrics\"])\r\n",
    "\r\n",
    "        # Update and save best result\r\n",
    "        if validation[\"metrics\"].loss < metrics_best.loss:\r\n",
    "            metrics_best.update(step, validation[\"metrics\"])\r\n",
    "            torch.save(model.state_dict(), best_path_model)\r\n",
    "\r\n",
    "        # write to summary writer\r\n",
    "        metrics_best.write(summary_writer, step)\r\n",
    "\r\n",
    "        epoch = step // steps_per_epoch\r\n",
    "        train_metrics_res = train[\"metrics\"].result(append_tag=False)\r\n",
    "        val_metrics_res = validation[\"metrics\"].result(append_tag=False)\r\n",
    "        metrics_strings = [\r\n",
    "            f\"{key}: train={train_metrics_res[key]:.6f}, val={val_metrics_res[key]:.6f}\"\r\n",
    "            for key in validation[\"metrics\"].keys\r\n",
    "        ]\r\n",
    "        logging.info(\r\n",
    "            f\"{step}/{num_steps} (epoch {epoch}): \" + \"; \".join(metrics_strings)\r\n",
    "        )\r\n",
    "\r\n",
    "        # decay learning rate on plateau\r\n",
    "        trainer.decay_maybe(validation[\"metrics\"].loss)\r\n",
    "\r\n",
    "        train[\"metrics\"].write(summary_writer, step)\r\n",
    "        validation[\"metrics\"].write(summary_writer, step)\r\n",
    "        train[\"metrics\"].reset_states()\r\n",
    "        validation[\"metrics\"].reset_states()\r\n",
    "\r\n",
    "        # Restore backup variables\r\n",
    "        trainer.restore_variable_backups()\r\n",
    "\r\n",
    "        # early stopping\r\n",
    "        if step - metrics_best.step > patience * evaluation_interval:\r\n",
    "            break\r\n",
    "\r\n",
    "result = {key + \"_best\": val for key, val in metrics_best.items()}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Print results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for key, val in metrics_best.items():\r\n",
    "    print(f\"{key}: {val}\")\r\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "6d9d58ddb04bb635eba824a3c64b6d0110bcc4c6cff8b192a6f7cbbb2bf10de4"
  },
  "kernelspec": {
   "display_name": "Python 3.5.4 64-bit",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": ""
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}