sbdd-controlled-diffusion.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Preliminary inputs\n",
    "pretrained_FDD_path = None # Insert the path to the pretrained FullDenoisingDiffusion model\n",
    "wandb_key = None # Insert your wandb code here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!git clone https://github.com/KaziFahmi/EDM_Custom"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append('/content/EDM_Custom')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install hydra-core torch_geometric wandb pytorch_lightning omegaconf torchmetrics setuptools scikit_learn imageio matplotlib numpy pandas Pillow tqdm rdkit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%cd EDM_Custom/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -r /content/EDM_Custom/requirements.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install rdkit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import torch\n",
    "import pytorch_lightning as pl\n",
    "\n",
    "import os\n",
    "import warnings\n",
    "import pathlib\n",
    "from pytorch_lightning import Trainer\n",
    "from pytorch_lightning.callbacks import ModelCheckpoint\n",
    "from pytorch_lightning.utilities.warnings import PossibleUserWarning\n",
    "\n",
    "from midi.datasets import qm9_dataset, geom_dataset\n",
    "from midi.diffusion_model import FullDenoisingDiffusion\n",
    "from controlnet.conditioned_diffusion_model import ConditionedDiffusion\n",
    "from omegaconf import OmegaConf\n",
    "from hydra import compose, initialize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "config_path = \"configs/\"  \n",
    "def load_hydra_configs(config_path):\n",
    "    with initialize(config_path=config_path):\n",
    "        cfg = compose(config_name=\"config\")\n",
    "    return cfg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\xcdtl\\AppData\\Local\\Temp\\ipykernel_16228\\1311385683.py:3: UserWarning: \n",
      "The version_base parameter is not specified.\n",
      "Please specify a compatability version level, or None.\n",
      "Will assume defaults for version 1.1\n",
      "  with initialize(config_path=config_path):\n"
     ]
    }
   ],
   "source": [
    "cfg = load_hydra_configs(config_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_config = cfg.dataset\n",
    "data_module = qm9_dataset.QM9DataModule(cfg)\n",
    "dataset_info = qm9_dataset.QM9infos(datamodule=data_module,cfg=cfg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pretrained_weights = torch.load(pretrained_FDD_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Marginal distribution of the classes: nodes: tensor([0.7231, 0.1151, 0.1591, 0.0026]) -- edges: tensor([0.7261, 0.2383, 0.0274, 0.0082, 0.0000]) -- charges: tensor([0.0078, 0.9705, 0.0218])\n"
     ]
    }
   ],
   "source": [
    "pretrained_FDD = FullDenoisingDiffusion(cfg=cfg, dataset_infos=dataset_info, train_smiles=[])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pretrained_FDD.load_state_dict(pretrained_weights['state_dict'])\n",
    "pretrained_EDM = pretrained_FDD.model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = ConditionedDiffusion(cft=cfg, dataset_infos=dataset_info, train_smiles=[], pretrained_EDM=pretrained_EDM)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "params_to_ignore = ['module.model.train_smiles', 'module.model.dataset_infos']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.nn.parallel.DistributedDataParallel._set_params_and_buffers_to_ignore_for_model(model, params_to_ignore)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint_callback = ModelCheckpoint(dirpath=f\"checkpoints/{cfg.general.name}\",\n",
    "                                              filename='{epoch}',\n",
    "                                              monitor='val/epoch_NLL',\n",
    "                                              save_top_k=5,\n",
    "                                              mode='min',\n",
    "                                              every_n_epochs=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "callbacks=[]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "last_ckpt_save = ModelCheckpoint(dirpath=f\"checkpoints/{cfg.general.name}\", filename='last', every_n_epochs=1)\n",
    "callbacks.append(checkpoint_callback)\n",
    "callbacks.append(last_ckpt_save)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "False"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import wandb\n",
    "wandb.login(key=wandb_key)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: False, used: False\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n",
      "c:\\Users\\xcdtl\\anaconda3\\envs\\midi\\lib\\site-packages\\pytorch_lightning\\trainer\\connectors\\logger_connector\\logger_connector.py:67: UserWarning: Starting from v1.9.0, `tensorboardX` has been removed as a dependency of the `pytorch_lightning` package, due to potential conflicts with other packages in the ML ecosystem. For this reason, `logger=True` will use `CSVLogger` as the default logger, unless the `tensorboard` or `tensorboardX` packages are found. Please `pip install lightning[extra]` or one of them to enable TensorBoard support by default\n",
      "  warning_cache.warn(\n"
     ]
    }
   ],
   "source": [
    "trainer = Trainer(gradient_clip_val=cfg.train.clip_grad, # Needed to load old checkpoints\n",
    "                      accelerator='gpu' if torch.cuda.is_available() else 'cpu',\n",
    "                      devices=cfg.general.gpus,\n",
    "                      max_epochs=cfg.train.n_epochs,\n",
    "                      check_val_every_n_epoch=cfg.general.check_val_every_n_epochs,\n",
    "                      fast_dev_run=cfg.general.name == 'debug',\n",
    "                      enable_progress_bar=cfg.train.progress_bar,\n",
    "                      callbacks=callbacks,\n",
    "                      log_every_n_steps=50,\n",
    "                      )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# FILEPATH: /f:/FYDP/EDM_Custom/colab_notebook.ipynb\n",
    "trainer.fit(model, datamodule=data_module)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tracemalloc\n",
    "\n",
    "def example_function():\n",
    "    tracemalloc.start()\n",
    "    # your code to profile\n",
    "    a = [1] * (10 ** 6)  # Example line to profile\n",
    "    snapshot = tracemalloc.take_snapshot()\n",
    "    top_stats = snapshot.statistics('lineno')\n",
    "\n",
    "    print(\"[ Top 10 ]\")\n",
    "    for stat in top_stats[:10]:\n",
    "        print(stat)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "midi",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}