example_training_LVAE_muSplit.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import os\n",
    "import socket\n",
    "import sys\n",
    "from copy import deepcopy\n",
    "from dataclasses import asdict\n",
    "from pathlib import Path\n",
    "from typing import Literal, Optional, Union\n",
    "\n",
    "import ml_collections\n",
    "import torch\n",
    "import wandb\n",
    "from pydantic import BaseModel, ConfigDict\n",
    "from pytorch_lightning import Trainer\n",
    "from pytorch_lightning.callbacks import (\n",
    "    EarlyStopping,\n",
    "    LearningRateMonitor,\n",
    "    ModelCheckpoint,\n",
    ")\n",
    "from pytorch_lightning.loggers import WandbLogger\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "\n",
    "from careamics.config import VAEAlgorithmConfig\n",
    "from careamics.config.architectures import LVAEModel\n",
    "from careamics.config.likelihood_model import (\n",
    "    GaussianLikelihoodConfig,\n",
    "    NMLikelihoodConfig,\n",
    ")\n",
    "from careamics.config.nm_model import GaussianMixtureNMConfig, MultiChannelNMConfig\n",
    "from careamics.config.optimizer_models import LrSchedulerModel, OptimizerModel\n",
    "from careamics.lightning import VAEModule\n",
    "from careamics.models.lvae.noise_models import noise_model_factory\n",
    "# TODO: sorry for this :(\n",
    "sys.path.insert(0, \"/home/federico.carrara/Documents/projects/microSplit-reproducibility/\")\n",
    "from data import (\n",
    "    LCMultiChDloader, MultiChDloader, DataSplitType, DataType\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set some parameters for the current training simulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "img_size: int = 64\n",
    "\"\"\"Spatial size of the input image.\"\"\"\n",
    "target_channels: int = 2\n",
    "\"\"\"Number of channels in the target image.\"\"\"\n",
    "multiscale_count: int = 3\n",
    "\"\"\"The number of LC inputs plus one (the actual input).\"\"\"\n",
    "predict_logvar: Optional[Literal[\"pixelwise\"]] = \"pixelwise\"\n",
    "\"\"\"Whether to compute also the log-variance as LVAE output.\"\"\"\n",
    "loss_type: Optional[Literal[\"musplit\", \"denoisplit\", \"denoisplit_musplit\"]] = \"musplit\"\n",
    "\"\"\"The type of reconstruction loss (i.e., likelihood) to use.\"\"\"\n",
    "nm_paths: Optional[tuple[str]] = [\n",
    "    \"/group/jug/ashesh/training_pre_eccv/noise_model/2402/221/GMMNoiseModel_ER-GT_all.mrc__6_4_Clip0.0-1.0_Sig0.125_UpNone_Norm0_bootstrap.npz\",\n",
    "    \"/group/jug/ashesh/training_pre_eccv/noise_model/2402/225/GMMNoiseModel_Microtubules-GT_all.mrc__6_4_Clip0.0-1.0_Sig0.125_UpNone_Norm0_bootstrap.npz\"\n",
    "]\n",
    "\"\"\"The paths to the pre-trained noise models for the different channels.\"\"\"\n",
    "# TODO: add denoisplit-musplit weights"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Other training parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: replace this with `careamics,config.training_model.TrainingConfig` once updated\n",
    "class TrainingConfig(BaseModel):\n",
    "    \"\"\"Configuration for training a VAE model.\"\"\"\n",
    "\n",
    "    model_config = ConfigDict(\n",
    "        validate_assignment=True, arbitrary_types_allowed=True, extra=\"allow\"\n",
    "    )\n",
    "\n",
    "    batch_size: int = 32\n",
    "    \"\"\"The batch size for training.\"\"\"\n",
    "    precision: int = 16\n",
    "    \"\"\"The precision to use for training.\"\"\"\n",
    "    lr: float = 1e-3\n",
    "    \"\"\"The learning rate for training.\"\"\"\n",
    "    lr_scheduler_patience: int = 30\n",
    "    \"\"\"The patience for the learning rate scheduler.\"\"\"\n",
    "    earlystop_patience: int = 200\n",
    "    \"\"\"The patience for the learning rate scheduler.\"\"\"\n",
    "    max_epochs: int = 400\n",
    "    \"\"\"The maximum number of epochs to train for.\"\"\"\n",
    "    num_workers: int = 4\n",
    "    \"\"\"The number of workers to use for data loading.\"\"\"\n",
    "    grad_clip_norm_value: int = 0.5\n",
    "    \"\"\"The value to use for gradient clipping (see lightning `Trainer`).\"\"\"\n",
    "    gradient_clip_algorithm: int = 'value'\n",
    "    \"\"\"The algorithm to use for gradient clipping (see lightning `Trainer`).\"\"\"\n",
    "\n",
    "train_config = TrainingConfig()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Create `Dataset` and `Dataloader`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get data config\n",
    "def get_data_config(\n",
    "):\n",
    "    data_config = ml_collections.ConfigDict()\n",
    "    data_config.data_dir = \"/group/jug/federico/careamics_training/data/BioSR\"\n",
    "    data_config.image_size = img_size #ok\n",
    "    data_config.target_channels = target_channels #ok\n",
    "    data_config.multiscale_lowres_count = multiscale_count #ok\n",
    "    data_config.data_type = DataType.BioSR_MRC #ok\n",
    "    data_config.ch1_fname = \"ER/GT_all.mrc\" #ok\n",
    "    data_config.ch2_fname = \"CCPs/GT_all.mrc\" #ok\n",
    "    data_config.poisson_noise_factor = -1 #ok\n",
    "    data_config.enable_gaussian_noise = True #ok\n",
    "    data_config.synthetic_gaussian_scale = 5100 #ok\n",
    "    data_config.input_has_dependant_noise = True #ok\n",
    "    return data_config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_dataset(\n",
    "    config: ml_collections.ConfigDict,\n",
    "    eval_datasplit_type = DataSplitType.Val,\n",
    "    skip_train_dataset = False,\n",
    "    kwargs_dict = None,\n",
    ") -> tuple[Dataset, Dataset, tuple[float, float]]:\n",
    "    if kwargs_dict is None:\n",
    "        kwargs_dict = {}\n",
    "\n",
    "    datapath = config.data_dir\n",
    "\n",
    "    # Hard-coded parameters (used to be in the config file)\n",
    "    normalized_input = True #ok\n",
    "    use_one_mu_std = True #ok\n",
    "    train_aug_rotate = False #ok\n",
    "    enable_random_cropping = True #ok\n",
    "    lowres_supervision = False #ok\n",
    "\n",
    "    # 1) Data loader for Lateral Contextualization\n",
    "    if config.multiscale_lowres_count > 1:\n",
    "        # Get padding attributes\n",
    "        if \"padding_kwargs\" not in kwargs_dict:\n",
    "            padding_kwargs = {\"mode\": \"reflect\"}\n",
    "        else:\n",
    "            padding_kwargs = kwargs_dict.pop(\"padding_kwargs\")\n",
    "\n",
    "        train_data = (\n",
    "            None\n",
    "            if skip_train_dataset\n",
    "            else LCMultiChDloader(\n",
    "                config,\n",
    "                datapath,\n",
    "                datasplit_type=DataSplitType.Train,\n",
    "                val_fraction=0.1,\n",
    "                test_fraction=0.1,\n",
    "                normalized_input=normalized_input,\n",
    "                use_one_mu_std=use_one_mu_std,\n",
    "                enable_rotation_aug=train_aug_rotate,\n",
    "                enable_random_cropping=enable_random_cropping,\n",
    "                num_scales=config.multiscale_lowres_count,\n",
    "                lowres_supervision=lowres_supervision,\n",
    "                padding_kwargs=padding_kwargs,\n",
    "                **kwargs_dict,\n",
    "                allow_generation=True,\n",
    "            )\n",
    "        )\n",
    "        max_val = train_data.get_max_val()\n",
    "\n",
    "        val_data = LCMultiChDloader(\n",
    "            config,\n",
    "            datapath,\n",
    "            datasplit_type=eval_datasplit_type,\n",
    "            val_fraction=0.1,\n",
    "            test_fraction=0.1,\n",
    "            normalized_input=normalized_input,\n",
    "            use_one_mu_std=use_one_mu_std,\n",
    "            enable_rotation_aug=False,  # No rotation aug on validation\n",
    "            enable_random_cropping=False,\n",
    "            # No random cropping on validation. Validation is evaluated on determistic grids\n",
    "            num_scales=config.multiscale_lowres_count,\n",
    "            lowres_supervision=lowres_supervision,\n",
    "            padding_kwargs=padding_kwargs,\n",
    "            allow_generation=False,\n",
    "            **kwargs_dict,\n",
    "            max_val=max_val,\n",
    "        )\n",
    "    # 2) Vanilla data loader\n",
    "    else:\n",
    "        train_data_kwargs = {\"allow_generation\": True, **kwargs_dict}\n",
    "        val_data_kwargs = {\"allow_generation\": False, **kwargs_dict}\n",
    "\n",
    "        train_data_kwargs[\"enable_random_cropping\"] = enable_random_cropping\n",
    "        val_data_kwargs[\"enable_random_cropping\"] = False\n",
    "\n",
    "        train_data = (\n",
    "            None\n",
    "            if skip_train_dataset\n",
    "            else MultiChDloader(\n",
    "                data_config=config,\n",
    "                fpath=datapath,\n",
    "                datasplit_type=DataSplitType.Train,\n",
    "                val_fraction=0.1,\n",
    "                test_fraction=0.1,\n",
    "                normalized_input=normalized_input,\n",
    "                use_one_mu_std=use_one_mu_std,\n",
    "                enable_rotation_aug=train_aug_rotate,\n",
    "                **train_data_kwargs,\n",
    "            )\n",
    "        )\n",
    "\n",
    "        max_val = train_data.get_max_val()\n",
    "        val_data = MultiChDloader(\n",
    "            data_config=config,\n",
    "            fpath=datapath,\n",
    "            datasplit_type=eval_datasplit_type,\n",
    "            val_fraction=0.1,\n",
    "            test_fraction=0.1,\n",
    "            normalized_input=normalized_input,\n",
    "            use_one_mu_std=use_one_mu_std,\n",
    "            enable_rotation_aug=False,  # No rotation aug on validation\n",
    "            max_val=max_val,\n",
    "            **val_data_kwargs,\n",
    "        )\n",
    "\n",
    "    mean_val, std_val = train_data.compute_mean_std()\n",
    "    train_data.set_mean_std(mean_val, std_val)\n",
    "    val_data.set_mean_std(mean_val, std_val)\n",
    "    data_stats = train_data.get_mean_std()\n",
    "\n",
    "    # NOTE: \"input\" mean & std are computed over the entire dataset and repeated for each channel.\n",
    "    # On the contrary, \"target\" mean & std are computed separately for each channel.\n",
    "    # manipulate data stats to only have one mean and std for the target\n",
    "    assert isinstance(data_stats, tuple)\n",
    "    assert isinstance(data_stats[0], dict)\n",
    "    data_stats = (\n",
    "        torch.tensor(data_stats[0][\"target\"]),\n",
    "        torch.tensor(data_stats[1][\"target\"])\n",
    "    )\n",
    "\n",
    "    return train_data, val_data, data_stats"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dset, val_dset, data_stats = create_dataset(\n",
    "    config=get_data_config(),\n",
    "    eval_datasplit_type=DataSplitType.Val,\n",
    "    skip_train_dataset=False,\n",
    "    kwargs_dict=None,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dloader = DataLoader(\n",
    "    train_dset,\n",
    "    batch_size=train_config.batch_size,\n",
    "    num_workers=train_config.num_workers,\n",
    "    shuffle=True\n",
    ")\n",
    "val_dloader = DataLoader(\n",
    "    val_dset,\n",
    "    batch_size=train_config.batch_size,\n",
    "    num_workers=train_config.num_workers,\n",
    "    shuffle=False\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Instantiate the lightning module"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_split_lightning_model(\n",
    "    algorithm: str,\n",
    "    loss_type: str,\n",
    "    img_size: int = 64,\n",
    "    multiscale_count: int = 1,\n",
    "    predict_logvar: Optional[Literal[\"pixelwise\"]] = None,\n",
    "    target_ch: int = 1,\n",
    "    NM_paths: Optional[list[Path]] = None,\n",
    "    training_config: TrainingConfig = TrainingConfig(),\n",
    "    data_mean: Optional[torch.Tensor] = None,\n",
    "    data_std: Optional[torch.Tensor] = None,\n",
    ") -> VAEModule:\n",
    "    \"\"\"Instantiate the muSplit lightining model.\"\"\"\n",
    "    lvae_config = LVAEModel(\n",
    "        architecture=\"LVAE\",\n",
    "        input_shape=img_size,\n",
    "        multiscale_count=multiscale_count,\n",
    "        z_dims=[128, 128, 128, 128],\n",
    "        output_channels=target_ch,\n",
    "        predict_logvar=predict_logvar,\n",
    "        analytical_kl=False, #ok\n",
    "    )\n",
    "\n",
    "    # gaussian likelihood\n",
    "    if loss_type in [\"musplit\", \"denoisplit_musplit\"]:\n",
    "        gaussian_lik_config = GaussianLikelihoodConfig(\n",
    "            predict_logvar=predict_logvar,\n",
    "            logvar_lowerbound=-5.,\n",
    "        )\n",
    "    else:\n",
    "        gaussian_lik_config = None\n",
    "    # noise model likelihood\n",
    "    if loss_type in [\"denoisplit\", \"denoisplit_musplit\"]:\n",
    "        assert NM_paths is not None, \"A path to a pre-trained noise model is required.\"\n",
    "        gmm_list = []\n",
    "        for NM_path in NM_paths:\n",
    "            gmm_list.append(\n",
    "                GaussianMixtureNMConfig(\n",
    "                    model_type=\"GaussianMixtureNoiseModel\",\n",
    "                    path=NM_path,\n",
    "                )\n",
    "            )\n",
    "        noise_model_config = MultiChannelNMConfig(noise_models=gmm_list)\n",
    "        nm = noise_model_factory(noise_model_config)\n",
    "        nm_lik_config = NMLikelihoodConfig(\n",
    "            noise_model=nm,\n",
    "            data_mean=data_mean,\n",
    "            data_std=data_std,\n",
    "        )\n",
    "    else:\n",
    "        noise_model_config = None\n",
    "        nm_lik_config = None\n",
    "\n",
    "    opt_config = OptimizerModel(\n",
    "        name=\"Adamax\",\n",
    "        parameters={\n",
    "            \"lr\": training_config.lr,\n",
    "            \"weight_decay\": 0,\n",
    "        },\n",
    "    )\n",
    "    lr_scheduler_config = LrSchedulerModel(\n",
    "        name=\"ReduceLROnPlateau\",\n",
    "        parameters={\n",
    "            \"mode\": \"min\",\n",
    "            \"factor\": 0.5,\n",
    "            \"patience\": training_config.lr_scheduler_patience,\n",
    "            \"verbose\": True,\n",
    "            \"min_lr\": 1e-12,\n",
    "        },\n",
    "    )\n",
    "\n",
    "    vae_config = VAEAlgorithmConfig(\n",
    "        algorithm_type=\"vae\",\n",
    "        algorithm=algorithm,\n",
    "        loss=loss_type,\n",
    "        model=lvae_config,\n",
    "        gaussian_likelihood_model=gaussian_lik_config,\n",
    "        noise_model=noise_model_config,\n",
    "        noise_model_likelihood_model=nm_lik_config,\n",
    "        optimizer=opt_config,\n",
    "        lr_scheduler=lr_scheduler_config,\n",
    "    )\n",
    "\n",
    "    return VAEModule(algorithm_config=vae_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "algo = \"musplit\" if loss_type == \"musplit\" else \"denoisplit\"\n",
    "lightning_model = create_split_lightning_model(\n",
    "    algorithm=algo,\n",
    "    loss_type=loss_type,\n",
    "    img_size=img_size,\n",
    "    multiscale_count=multiscale_count,\n",
    "    predict_logvar=predict_logvar,\n",
    "    target_ch=target_channels,\n",
    "    NM_paths=nm_paths,\n",
    "    training_config=train_config,\n",
    "    data_mean=data_stats[0],\n",
    "    data_std=data_stats[1],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Set utils for training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datetime import datetime\n",
    "\n",
    "from careamics.lvae_training.train_utils import get_new_model_version\n",
    "\n",
    "\n",
    "def get_new_model_version(model_dir: Union[Path, str]) -> int:\n",
    "    \"\"\"Create a unique version ID for a new model run.\"\"\"\n",
    "    versions = []\n",
    "    for version_dir in os.listdir(model_dir):\n",
    "        try:\n",
    "            versions.append(int(version_dir))\n",
    "        except:\n",
    "            print(\n",
    "                f\"Invalid subdirectory:{model_dir}/{version_dir}. Only integer versions are allowed\"\n",
    "            )\n",
    "            exit()\n",
    "    if len(versions) == 0:\n",
    "        return \"0\"\n",
    "    return f\"{max(versions) + 1}\"\n",
    "\n",
    "def get_workdir(\n",
    "    root_dir: str,\n",
    "    model_name: str,\n",
    ") -> tuple[Path, Path]:\n",
    "    \"\"\"Get the workdir for the current model.\n",
    "    \n",
    "    It has the following structure: \"root_dir/YYMM/model_name/version\"\n",
    "    \"\"\"\n",
    "    rel_path = datetime.now().strftime(\"%y%m\")\n",
    "    cur_workdir = os.path.join(root_dir, rel_path)\n",
    "    Path(cur_workdir).mkdir(exist_ok=True)\n",
    "\n",
    "    rel_path = os.path.join(rel_path, model_name)\n",
    "    cur_workdir = os.path.join(root_dir, rel_path)\n",
    "    Path(cur_workdir).mkdir(exist_ok=True)\n",
    "\n",
    "    rel_path = os.path.join(rel_path, get_new_model_version(cur_workdir))\n",
    "    cur_workdir = os.path.join(root_dir, rel_path)\n",
    "    try:\n",
    "        Path(cur_workdir).mkdir(exist_ok=False)\n",
    "    except FileExistsError:\n",
    "        print(\n",
    "            f\"Workdir {cur_workdir} already exists.\"\n",
    "        )\n",
    "    return cur_workdir, rel_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ROOT_DIR = \"/put/your/root/dir/here/where/to/store/model/ckpts/and/stuff\"\n",
    "lc_tag = \"with\" if multiscale_count > 1 else \"no\"\n",
    "workdir, exp_tag = get_workdir(ROOT_DIR, f\"{algo}_{lc_tag}_LC\")\n",
    "print(f\"Current workdir: {workdir}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the logger\n",
    "project_name = \"_\".join((\"careamics\", algo))\n",
    "if project_name == \"_\".join((\"careamics\", algo)):\n",
    "    raise ValueError(\"Please create your own project name for wandb.\")\n",
    "custom_logger = WandbLogger(\n",
    "    name=os.path.join(socket.gethostname(), exp_tag),\n",
    "    save_dir=workdir,\n",
    "    project=project_name,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define callbacks (e.g., ModelCheckpoint, EarlyStopping, etc.)\n",
    "custom_callbacks = [\n",
    "    EarlyStopping(\n",
    "        monitor=\"val_loss\",\n",
    "        min_delta=1e-6,\n",
    "        patience=train_config.earlystop_patience,\n",
    "        mode=\"min\",\n",
    "        verbose=True,\n",
    "    ),\n",
    "    ModelCheckpoint(\n",
    "        dirpath=workdir,\n",
    "        filename=\"best-{epoch}\",\n",
    "        monitor=\"val_loss\",\n",
    "        save_top_k=1,\n",
    "        save_last=True,\n",
    "        mode=\"min\",\n",
    "    ),\n",
    "    LearningRateMonitor(logging_interval=\"epoch\")\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Save configs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get configs\n",
    "algo_config = lightning_model.algorithm_config\n",
    "data_config = get_data_config()\n",
    "# temp -> remove fields that we don't want to save\n",
    "loss_config = deepcopy(asdict(lightning_model.loss_parameters))\n",
    "del loss_config[\"noise_model_likelihood\"]\n",
    "del loss_config[\"gaussian_likelihood\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save Configs JSON\n",
    "with open(os.path.join(workdir, \"algorithm_config.json\"), \"w\") as f:\n",
    "    f.write(algo_config.model_dump_json(indent=4))\n",
    "\n",
    "with open(os.path.join(workdir, \"training_config.json\"), \"w\") as f:\n",
    "    f.write(train_config.model_dump_json(indent=4))\n",
    "\n",
    "with open(os.path.join(workdir, \"data_config.json\"), \"w\") as f:\n",
    "    json.dump(data_config.to_dict(), f, indent=4)\n",
    "\n",
    "with open(os.path.join(workdir, \"loss_config.json\"), \"w\") as f:\n",
    "    json.dump(loss_config, f, indent=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save Configs in WANDB\n",
    "custom_logger.experiment.config.update({\n",
    "    \"algorithm\": algo_config.model_dump()\n",
    "})\n",
    "\n",
    "custom_logger.experiment.config.update({\n",
    "    \"training\": train_config.model_dump()\n",
    "})\n",
    "\n",
    "custom_logger.experiment.config.update({\n",
    "    \"data\": data_config.to_dict()\n",
    "})\n",
    "\n",
    "custom_logger.experiment.config.update({\n",
    "        \"loss_params\": loss_config\n",
    "})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Train the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = Trainer(\n",
    "    max_epochs=train_config.max_epochs,\n",
    "    accelerator=\"gpu\",\n",
    "    enable_progress_bar=True,\n",
    "    logger=custom_logger,\n",
    "    callbacks=custom_callbacks,\n",
    "    precision=train_config.precision,\n",
    "    gradient_clip_val=train_config.grad_clip_norm_value, # only works with `accelerator=\"gpu\"`\n",
    "    gradient_clip_algorithm=train_config.gradient_clip_algorithm,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# os.environ[\"WANDB_MODE\"] = \"disabled\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.fit(\n",
    "    model=lightning_model,\n",
    "    train_dataloaders=train_dloader,\n",
    "    val_dataloaders=val_dloader,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wandb.finish()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "train_lvae",
   "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.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}