execute.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/upriyam-cmu/EDGE-Rec/blob/main/execute.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "id": "9Q6JZwINFIas"
   },
   "source": [
    "# Install project code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "id": "T8zKX7jqFIax"
   },
   "outputs": [],
   "source": [
    "!pip uninstall -y edge-rec\n",
    "!pip install -e ."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "id": "YyB0la9eFIa2"
   },
   "source": [
    "# Train model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# from edge_rec.utils import load_config\n",
    "# data_holder, model, trainer = load_config(\"./config/latest.json\")"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "qHy9cdt2FIa3"
   },
   "outputs": [],
   "source": [
    "from edge_rec.datasets import MovieLensDataHolder, RatingsTransform, FeatureTransform\n",
    "\n",
    "from edge_rec.model import GraphReconstructionModel, GraphTransformer\n",
    "from edge_rec.model.embed import MovieLensFeatureEmbedder, SinusoidalPositionalEmbedding\n",
    "\n",
    "from edge_rec.diffusion import GaussianDiffusion\n",
    "from edge_rec.exec import Trainer, compute_metrics_from_ratings as compute_metrics\n",
    "\n",
    "from torch import nn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "data_holder = MovieLensDataHolder(\n",
    "    ml100k=True,\n",
    "    augmentations=dict(\n",
    "        ratings=RatingsTransform.ToGaussian(),\n",
    "        rating_counts=FeatureTransform.LogPolynomial(2),  # degree 2 --> dim_size = 2 (for embedder, below)\n",
    "    ),\n",
    ")"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "embed = MovieLensFeatureEmbedder(\n",
    "    ml100k=True,\n",
    "    user_id_dim=128,\n",
    "    user_age_dim=None,\n",
    "    user_gender_dim=None,\n",
    "    user_occupation_dim=None,\n",
    "    user_rating_counts_dims=None,\n",
    "    movie_id_dim=128,\n",
    "    movie_genre_ids_dim=None,\n",
    "    movie_genre_multihot_dims=None,\n",
    "    movie_rating_counts_dims=None,\n",
    ")\n",
    "core = GraphTransformer(\n",
    "    n_blocks=16,\n",
    "    n_channels=1,\n",
    "    n_channels_internal=5,\n",
    "    n_features=embed.output_sizes,\n",
    "    time_embedder=SinusoidalPositionalEmbedding(32),\n",
    "    attn_kwargs=dict(heads=4, dim_head=32, num_mem_kv=4, speed_hack=True, share_weights=False, dropout=0.1),\n",
    "    feed_forward_kwargs=dict(hidden_dims=(2, 4, 2), activation_fn=\"nn.SiLU\"),\n",
    ")\n",
    "model = GraphReconstructionModel(embed, core, feature_dim_size=None)\n",
    "\n",
    "print(\"model size:\", model.model_size)\n",
    "print(\"embedding size:\", model.embedding.model_size)\n",
    "print(\"transformer size:\", model.core_model.model_size)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "diffusion_model = GaussianDiffusion(\n",
    "    model=model,\n",
    "    image_size=50,\n",
    "    offset_noise_strength=0.1,\n",
    "    p_losses_weight=1.,\n",
    "    bayes_personalized_ranking_loss_weight=0.,\n",
    ")\n",
    "trainer = Trainer(\n",
    "    # model\n",
    "    diffusion_model=diffusion_model,\n",
    "    # datasets\n",
    "    data_holder=data_holder,\n",
    "    subgraph_size=50,\n",
    "    target_density=None,\n",
    "    # training\n",
    "    batch_size=1,\n",
    "    gradient_accumulate_every=1,\n",
    "    force_batch_size=True,\n",
    "    train_num_steps=int(1e5),\n",
    "    train_mask_unknown_ratings=True,\n",
    "    # eval\n",
    "    eval_batch_size=None,  # copy training batch size if None\n",
    "    n_eval_iters=100,\n",
    "    eval_every=200,\n",
    "    sample_on_eval=False,\n",
    "    # optim\n",
    "    train_lr=1e-4,\n",
    "    adam_betas=(0.9, 0.99),\n",
    "    max_grad_norm=1.,\n",
    "    # logging\n",
    "    results_folder=\"./results\",\n",
    "    ema_update_every=10,\n",
    "    ema_decay=0.995,\n",
    "    save_every_nth_eval=1,\n",
    "    score_on_save=False,  # suggested: copy value from use_wandb\n",
    "    use_wandb=False,\n",
    "    save_config=True,\n",
    "    # accelerator\n",
    "    amp=False,\n",
    "    mixed_precision_type='fp16',\n",
    "    split_batches=True,\n",
    ")\n",
    "print(\"Using device:\", trainer.device)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "9fJB149xFIa5"
   },
   "outputs": [],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "id": "d5_rCKjIFIa6"
   },
   "source": [
    "# Sample ratings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "all_metrics = {\n",
    "    trainer.score(\n",
    "        milestone=milestone,\n",
    "        predicting_ratings=True,\n",
    "        n_samples=10,\n",
    "        do_inpainting_sampling=True,\n",
    "        silence_inner_tqdm=True,\n",
    "    )\n",
    "    for milestone in (4000, 6000, 8000, 10000)\n",
    "}"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "id": "uoaM5uiZFIa7"
   },
   "source": [
    "# Evaluate metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.ticker import (MultipleLocator, AutoMinorLocator)\n",
    "\n",
    "mpl.rcParams.update({'figure.dpi': 300, 'font.size': 4})\n",
    "plt.style.use('bmh')\n",
    "\n",
    "hrv_name = {\n",
    "    'precision': 'Precision',\n",
    "    'recall': 'Recall',\n",
    "    'mean_reciprocal_rank': 'MRR',\n",
    "    'hit_rate': 'HR',\n",
    "    'ndcg': 'NDCG'\n",
    "}\n",
    "\n",
    "\n",
    "def plot_metrics(metrics: dict, plot_name: str, nested_key_name=None):\n",
    "    def _format_plot(ax, curve_name, min_value, max_value):\n",
    "        ax.set_title(hrv_name[curve_name])\n",
    "        ax.xaxis.grid(True, which='major')\n",
    "        ax.xaxis.set_major_locator(MultipleLocator(10))\n",
    "        ax.xaxis.set_major_formatter('{x:.0f}')\n",
    "\n",
    "        ax.set_xlabel(\"Top-K\")\n",
    "        ax.set_aspect(50 / (max_value - min_value))\n",
    "\n",
    "    def _plot(ax, data):\n",
    "        ks = (1, 5, 10, 20, 30, 40, 50)\n",
    "        ax.plot(ks, data, '.-', linewidth=1.0, markersize=4.0)\n",
    "\n",
    "    if nested_key_name is not None:\n",
    "        keys = sorted(list(metrics.keys()))\n",
    "    else:\n",
    "        keys = None\n",
    "\n",
    "    fig, axs = plt.subplots(1, 5)\n",
    "    for ax, name in zip(axs, ['precision', 'recall', 'ndcg', 'mean_reciprocal_rank', 'hit_rate']):\n",
    "        if nested_key_name is None:\n",
    "            _format_plot(ax, name, min_value=np.min(metrics[name]), max_value=np.max(metrics[name]))\n",
    "            _plot(ax, metrics[name])\n",
    "        else:\n",
    "            all_data = [m[name] for m in metrics.values()]\n",
    "            _format_plot(ax, name, min_value=np.min(all_data), max_value=np.max(all_data))\n",
    "            for key in keys:\n",
    "                _plot(ax, metrics[key][name])\n",
    "\n",
    "        if name == 'precision':\n",
    "            ax.set_ylabel(plot_name)\n",
    "\n",
    "    plt.tight_layout(h_pad=-25.0)\n",
    "    if nested_key_name is not None:\n",
    "        plt.legend([f\"{nested_key_name}={key}\" for key in keys])\n",
    "\n",
    "    plt.show()"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "plot_metrics(metrics=all_metrics, plot_name=f\"ML-100k\", nested_key_name=\"milestone\")"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "id": "RpzXY9IfFIa-"
   },
   "source": [
    "# Display sampled ratings distribution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "denoised_graph = trainer.eval(\n",
    "    milestone=10000,\n",
    "    do_inpainting_sampling=True,\n",
    "    silence_inner_tqdm=True,\n",
    "    save_sampled_graph=False,\n",
    ")"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "37uiTgF2FIa_"
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(3, 2))\n",
    "plt.hist(data_holder.ratings_transform.invert(denoised_graph).numpy().flatten(), bins=20)\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "include_colab_link": true,
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "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.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}