models/toy_sources/toy_sources.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d-vhj52HLOR5"
      },
      "source": [
        "# Rate–distortion experiments with toy sources\n",
        "\n",
        "This notebook contains code to train VECVQ and NTC models using stochastic rate–distortion optimization.\n",
        "\n",
        "The Laplace and Banana sources are described in:\n",
        "\n",
        "> \"Nonlinear Transform Coding\"<br />\n",
        "> J. Ballé, P. A. Chou, D. Minnen, S. Singh, N. Johnston, E. Agustsson, S. J. Hwang, G. Toderici<br />\n",
        "> https://arxiv.org/abs/2007.03034\n",
        "\n",
        "The Sawbridge process is described in:\n",
        "\n",
        "> \"Neural Networks Optimally Compress the Sawbridge\"<br />\n",
        "> A. B. Wagner, J. Ballé<br />\n",
        "> https://arxiv.org/abs/2011.05065\n",
        "\n",
        "This notebook requires TFC v2 (`pip install tensorflow-compression==2.*`)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "wdA5NUZ-fSxG"
      },
      "outputs": [],
      "source": [
        "#@title Dependencies for Colab\n",
        "\n",
        "# Run this cell to install the necessary dependencies when running the notebook\n",
        "# directly in a Colaboratory hosted runtime from Github.\n",
        "\n",
        "# Installs the latest version of TFC compatible with the installed TF version.\n",
        "!pip install tensorflow-compression~=$(pip show tensorflow | perl -p -0777 -e 's/.*Version: (\\d+\\.\\d+).*/\\1.0/sg')\n",
        "\n",
        "# Downloads the 'models' directory from Github.\n",
        "![[ -e /tfc ]] || git clone https://github.com/tensorflow/compression /tfc\n",
        "%cd /tfc/models\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "6d0oYddBWybj"
      },
      "outputs": [],
      "source": [
        "#@title Imports\n",
        "\n",
        "from absl import logging\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "import tensorflow_compression as tfc\n",
        "import tensorflow_probability as tfp\n",
        "\n",
        "tfm = tf.math\n",
        "tfkl = tf.keras.layers\n",
        "tfpb = tfp.bijectors\n",
        "tfpd = tfp.distributions\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "T67H_rSSWGmh"
      },
      "outputs": [],
      "source": [
        "#@title Matplotlib configuration\n",
        "\n",
        "import cycler\n",
        "import matplotlib as mpl\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "_colors = [\n",
        "    \"#1f77b4\", \"#ff7f0e\", \"#2ca02c\", \"#d62728\", \"#9467bd\",\n",
        "    \"#8c564b\", \"#e377c2\", \"#7f7f7f\", \"#bcbd22\", \"#17becf\",\n",
        "]\n",
        "\n",
        "plt.rc(\"axes\", facecolor=\"white\", labelsize=\"large\",\n",
        "       prop_cycle=cycler.cycler(color=_colors))\n",
        "plt.rc(\"grid\", color=\"black\", alpha=.1)\n",
        "plt.rc(\"legend\", frameon=True, framealpha=.9, borderpad=.5, handleheight=1,\n",
        "       fontsize=\"large\")\n",
        "plt.rc(\"image\", cmap=\"viridis\", interpolation=\"nearest\")\n",
        "plt.rc(\"figure\", figsize=(16, 8))\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "vawgWpGxXHIj"
      },
      "outputs": [],
      "source": [
        "#@title Source definitions\n",
        "\n",
        "from toy_sources import sawbridge\n",
        "\n",
        "\n",
        "def _rotation_2d(degrees):\n",
        "  phi = tf.convert_to_tensor(degrees / 180 * np.pi, dtype=tf.float32)\n",
        "  rotation = [[tfm.cos(phi), -tfm.sin(phi)], [tfm.sin(phi), tfm.cos(phi)]]\n",
        "  rotation = tf.linalg.LinearOperatorFullMatrix(\n",
        "      rotation, is_non_singular=True, is_square=True)\n",
        "  return rotation\n",
        "\n",
        "\n",
        "def get_laplace(loc=0, scale=1):\n",
        "  return tfpd.Independent(\n",
        "      tfpd.Laplace(loc=[loc], scale=[scale]),\n",
        "      reinterpreted_batch_ndims=1,\n",
        "  )\n",
        "\n",
        "\n",
        "def get_banana():\n",
        "  return tfpd.TransformedDistribution(\n",
        "      tfpd.Independent(tfpd.Normal(loc=[0, 0], scale=[3, .5]), 1),\n",
        "      tfpb.Invert(tfpb.Chain([\n",
        "          tfpb.RealNVP(\n",
        "              num_masked=1,\n",
        "              shift_and_log_scale_fn=lambda x, _: (.1 * x ** 2, None)),\n",
        "          tfpb.ScaleMatvecLinearOperator(_rotation_2d(240)),\n",
        "          tfpb.Shift([1, 1]),\n",
        "      ])),\n",
        "  )\n",
        "\n",
        "\n",
        "def get_sawbridge(order=1, stationary=False, num_points=1024):\n",
        "  index_points = tf.linspace(0., 1., num_points)\n",
        "  return sawbridge.Sawbridge(\n",
        "      index_points, stationary=stationary, order=order)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "BFbkxNBXXKSX"
      },
      "outputs": [],
      "source": [
        "#@title Model definitions\n",
        "\n",
        "import pywt\n",
        "import scipy\n",
        "from toy_sources import ntc\n",
        "from toy_sources import vecvq\n",
        "\n",
        "\n",
        "def _get_activation(activation, dtype):\n",
        "  if not activation:\n",
        "    return None\n",
        "  if activation == \"gdn\":\n",
        "    return tfc.GDN(dtype=dtype)\n",
        "  elif activation == \"igdn\":\n",
        "    return tfc.GDN(inverse=True, dtype=dtype)\n",
        "  else:\n",
        "    return getattr(tf.nn, activation)\n",
        "\n",
        "\n",
        "def _make_nlp(units, activation, name, input_shape, dtype):\n",
        "  kwargs = [dict(  # pylint:disable=g-complex-comprehension\n",
        "      units=u, use_bias=True, activation=activation,\n",
        "      name=f\"{name}_{i}\", dtype=dtype,\n",
        "  ) for i, u in enumerate(units)]\n",
        "  kwargs[0].update(input_shape=input_shape)\n",
        "  kwargs[-1].update(activation=None)\n",
        "  return tf.keras.Sequential(\n",
        "      [tf.keras.layers.Dense(**k) for k in kwargs], name=name)\n",
        "\n",
        "\n",
        "def get_ntc_mlp_model(analysis_filters, synthesis_filters,\n",
        "                      analysis_activation, synthesis_activation,\n",
        "                      latent_dims, source, dtype=tf.float32, **kwargs):\n",
        "  \"\"\"NTC with MLP transforms.\"\"\"\n",
        "  source_dims, = source.event_shape\n",
        "\n",
        "  analysis = _make_nlp(\n",
        "      analysis_filters + [latent_dims],\n",
        "      _get_activation(analysis_activation, dtype),\n",
        "      \"analysis\",\n",
        "      [source_dims],\n",
        "      dtype,\n",
        "  )\n",
        "  synthesis = _make_nlp(\n",
        "      synthesis_filters + [source_dims],\n",
        "      _get_activation(synthesis_activation, dtype),\n",
        "      \"synthesis\",\n",
        "      [latent_dims],\n",
        "      dtype,\n",
        "  )\n",
        "\n",
        "  return ntc.NTCModel(\n",
        "      analysis=analysis, synthesis=synthesis, source=source, dtype=dtype,\n",
        "      **kwargs)\n",
        "\n",
        "\n",
        "def get_ltc_model(latent_dims, source, dtype=tf.float32, **kwargs):\n",
        "  \"\"\"LTC.\"\"\"\n",
        "  source_dims, = source.event_shape\n",
        "\n",
        "  analysis = tf.keras.Sequential([\n",
        "      tf.keras.layers.Dense(\n",
        "          latent_dims, use_bias=True, activation=None, name=\"analysis\",\n",
        "          input_shape=[source_dims], dtype=dtype),\n",
        "  ], name=\"analysis\")\n",
        "  synthesis = tf.keras.Sequential([\n",
        "      tf.keras.layers.Dense(\n",
        "          source_dims, use_bias=True, activation=None, name=\"synthesis\",\n",
        "          input_shape=[latent_dims], dtype=dtype),\n",
        "  ], name=\"synthesis\")\n",
        "\n",
        "  return ntc.NTCModel(\n",
        "      analysis=analysis, synthesis=synthesis, source=source, dtype=dtype,\n",
        "      **kwargs)\n",
        "\n",
        "\n",
        "@tf.function\n",
        "def estimate_klt(source, num_samples, latent_dims):\n",
        "  \"\"\"Estimates KLT.\"\"\"\n",
        "\n",
        "  dims = source.event_shape[0]\n",
        "\n",
        "  def energy(samples):\n",
        "    c = tf.linalg.matmul(samples, samples, transpose_a=True)\n",
        "    return c / tf.cast(num_samples[1], tf.float32)\n",
        "\n",
        "  # Estimate mean.\n",
        "  mean = tf.zeros([dims])\n",
        "  for _ in range(num_samples[0]):\n",
        "    samples = source.sample(num_samples[1])\n",
        "    mean += tf.reduce_mean(samples, axis=0)\n",
        "  mean /= tf.cast(num_samples[0], tf.float32)\n",
        "\n",
        "  # Estimate covariance.\n",
        "  covariance = tf.zeros([dims, dims])\n",
        "  for _ in range(num_samples[0]):\n",
        "    samples = source.sample(num_samples[1])\n",
        "    covariance += energy(samples - mean)\n",
        "  covariance /= tf.cast(num_samples[0], tf.float32)\n",
        "\n",
        "  variance = tf.reduce_sum(tf.linalg.diag_part(covariance))\n",
        "  tf.print(\"SOURCE VARIANCE:\", variance)\n",
        "\n",
        "  # Compute first latent_dims eigenvalues in descending order.\n",
        "  eig, eigv = tf.linalg.eigh(covariance)\n",
        "  eig = eig[::-1]\n",
        "  eigv = eigv[:, ::-1]\n",
        "  eig = eig[:latent_dims]\n",
        "  eigv = eigv[:, :latent_dims]\n",
        "  tf.print(\"SOURCE EIGENVALUES:\", eig)\n",
        "\n",
        "  # Estimate covariance again after whitening.\n",
        "  whitened = tf.zeros([latent_dims, latent_dims])\n",
        "  for _ in range(num_samples[0]):\n",
        "    samples = source.sample(num_samples[1])\n",
        "    whitened += energy(tf.linalg.matmul(samples - mean, eigv))\n",
        "  whitened /= tf.cast(num_samples[0], tf.float32)\n",
        "  whitened_var = tf.linalg.diag_part(whitened)\n",
        "  whitened /= tf.sqrt(\n",
        "      whitened_var[:, None] * whitened_var[None, :]) + 1e-20\n",
        "  error = tf.linalg.set_diag(abs(whitened), tf.zeros(latent_dims))\n",
        "  error = tf.reduce_max(error)\n",
        "  tf.print(\"MAX. CORRELATION COEFFICIENT:\", error)\n",
        "\n",
        "  return eigv, error\n",
        "\n",
        "\n",
        "class ScaleAndBias(tf.keras.layers.Layer):\n",
        "  \"\"\"Multiplies each channel by a learned scaling factor and adds a bias.\"\"\"\n",
        "\n",
        "  def __init__(self, scale_first, init_scale=1, **kwargs):\n",
        "    super().__init__(**kwargs)\n",
        "    self.scale_first = bool(scale_first)\n",
        "    self.init_scale = float(init_scale)\n",
        "\n",
        "  def build(self, input_shape):\n",
        "    input_shape = tf.TensorShape(input_shape)\n",
        "    channels = int(input_shape[-1])\n",
        "    self._log_factors = self.add_weight(\n",
        "        name=\"log_factors\", shape=[channels],\n",
        "        initializer=tf.keras.initializers.Constant(\n",
        "            tf.math.log(self.init_scale)))\n",
        "    self.bias = self.add_weight(\n",
        "        name=\"bias\", shape=[channels],\n",
        "        initializer=tf.keras.initializers.Zeros())\n",
        "    super().build(input_shape)\n",
        "\n",
        "  @property\n",
        "  def factors(self):\n",
        "    return tf.math.exp(self._log_factors)\n",
        "\n",
        "  def call(self, inputs):\n",
        "    if self.scale_first:\n",
        "      return inputs * self.factors + self.bias\n",
        "    else:\n",
        "      return (inputs + self.bias) * self.factors\n",
        "\n",
        "\n",
        "def get_ltc_klt_model(latent_dims, source, num_samples, tolerance, \n",
        "                      dtype=tf.float32, **kwargs):\n",
        "  \"\"\"LTC constrained to KLT.\"\"\"\n",
        "  source_dims, = source.event_shape\n",
        "\n",
        "  # Estimate KLT from samples.\n",
        "  eigv, error = estimate_klt(\n",
        "      source, tf.constant(num_samples), tf.constant(latent_dims))\n",
        "  assert error < tolerance, error.numpy()\n",
        "  eigv = tf.cast(eigv, dtype)\n",
        "\n",
        "  analysis = tf.keras.Sequential([\n",
        "      tf.keras.layers.Dense(\n",
        "          latent_dims, use_bias=False, activation=None, name=\"klt\",\n",
        "          kernel_initializer=lambda *a, **k: eigv,\n",
        "          trainable=False, input_shape=[source_dims], dtype=dtype),\n",
        "      ScaleAndBias(\n",
        "          scale_first=True, name=\"klt_scaling\", dtype=dtype),\n",
        "  ], name=\"analysis\")\n",
        "  synthesis = tf.keras.Sequential([\n",
        "      ScaleAndBias(\n",
        "          scale_first=False, name=\"iklt_scaling\",\n",
        "          input_shape=[latent_dims], dtype=dtype),\n",
        "      tf.keras.layers.Dense(\n",
        "          source_dims, use_bias=False, activation=None, name=\"iklt\",\n",
        "          kernel_initializer=lambda *a, **k: tf.transpose(eigv),\n",
        "          trainable=False, dtype=dtype),\n",
        "  ], name=\"synthesis\")\n",
        "\n",
        "  return ntc.NTCModel(\n",
        "      analysis=analysis, synthesis=synthesis, source=source, dtype=dtype,\n",
        "      **kwargs)\n",
        "\n",
        "\n",
        "def get_ltc_ortho_model(latent_dims, source, transform, dtype=tf.float32,\n",
        "                        **kwargs):\n",
        "  \"\"\"LTC constrained to fixed orthonormal transforms.\"\"\"\n",
        "  source_dims, = source.event_shape\n",
        "\n",
        "  if transform == \"dct\":\n",
        "    basis = scipy.fftpack.dct(np.eye(source_dims), norm=\"ortho\")\n",
        "  else:\n",
        "    num_levels = int(round(np.log2(source_dims)))\n",
        "    assert 2 ** num_levels == source_dims\n",
        "    basis = []\n",
        "    for impulse in np.eye(source_dims):\n",
        "      levels = pywt.wavedec(\n",
        "          impulse, transform, mode=\"periodization\", level=num_levels)\n",
        "      basis.append(np.concatenate(levels))\n",
        "    basis = np.array(basis)\n",
        "\n",
        "  # `basis` must have IO format, so DC should be in first column.\n",
        "  assert np.allclose(basis[:, 0], basis[0, 0])\n",
        "  assert not np.allclose(basis[0, :], basis[0, 0])\n",
        "\n",
        "  # `basis` should be orthonormal.\n",
        "  assert np.allclose(np.dot(basis, basis.T), np.eye(source_dims))\n",
        "\n",
        "  # Only take the first `latent_dims` basis functions.\n",
        "  basis = tf.constant(basis[:, :latent_dims], dtype=dtype)\n",
        "\n",
        "  analysis = tf.keras.Sequential([\n",
        "      tf.keras.layers.Dense(\n",
        "          latent_dims, use_bias=False, activation=None, name=transform,\n",
        "          kernel_initializer=lambda *a, **k: basis,\n",
        "          trainable=False, input_shape=[source_dims], dtype=dtype),\n",
        "      ScaleAndBias(\n",
        "          scale_first=True, name=f\"{transform}_scaling\", dtype=dtype),\n",
        "  ], name=\"analysis\")\n",
        "  synthesis = tf.keras.Sequential([\n",
        "      ScaleAndBias(\n",
        "          scale_first=False, name=f\"i{transform}_scaling\",\n",
        "          input_shape=[latent_dims], dtype=dtype),\n",
        "      tf.keras.layers.Dense(\n",
        "          source_dims, use_bias=False, activation=None, name=f\"i{transform}\",\n",
        "          kernel_initializer=lambda *a, **k: tf.transpose(basis),\n",
        "          trainable=False, dtype=dtype),\n",
        "  ], name=\"synthesis\")\n",
        "\n",
        "  return ntc.NTCModel(\n",
        "      analysis=analysis, synthesis=synthesis, source=source, dtype=dtype,\n",
        "      **kwargs)\n",
        "\n",
        "\n",
        "def get_vecvq_model(**kwargs):\n",
        "  return vecvq.VECVQModel(**kwargs)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "Sx2-WKHQZbUY"
      },
      "outputs": [],
      "source": [
        "#@title Learning schedule definitions\n",
        "\n",
        "def get_lr_scheduler(learning_rate, epochs, warmup_epochs=0):\n",
        "  \"\"\"Returns a learning rate scheduler function for the given configuration.\"\"\"\n",
        "  def scheduler(epoch, lr):\n",
        "    del lr  # unused\n",
        "    if epoch < warmup_epochs:\n",
        "      return learning_rate * 10. ** (epoch - warmup_epochs)\n",
        "    if epoch < 1/2 * epochs:\n",
        "      return learning_rate\n",
        "    if epoch < 3/4 * epochs:\n",
        "      return learning_rate * 1e-1\n",
        "    if epoch < 7/8 * epochs:\n",
        "      return learning_rate * 1e-2\n",
        "    return learning_rate * 1e-3\n",
        "  return scheduler\n",
        "\n",
        "\n",
        "class AlphaScheduler(tf.keras.callbacks.Callback):\n",
        "  \"\"\"Alpha parameter scheduler.\"\"\"\n",
        "\n",
        "  def __init__(self, schedule, verbose=0):\n",
        "    super().__init__()\n",
        "    self.schedule = schedule\n",
        "    self.verbose = verbose\n",
        "\n",
        "  def on_epoch_begin(self, epoch, logs=None):\n",
        "    if not hasattr(self.model, \"alpha\"):\n",
        "      # Silently ignore models that don't have an alpha parameter.\n",
        "      return\n",
        "    self.model.force_alpha = self.schedule(epoch)\n",
        "\n",
        "  def on_epoch_end(self, epoch, logs=None):\n",
        "    if not hasattr(self.model, \"alpha\"):\n",
        "      # Silently ignore models that don't have an alpha parameter.\n",
        "      return\n",
        "    if not hasattr(self.model, \"soft_round\") or not any(self.model.soft_round):\n",
        "      # Silently ignore models that don't use soft rounding.\n",
        "      return\n",
        "    logs[\"alpha\"] = self.model.alpha\n",
        "\n",
        "\n",
        "def get_alpha_scheduler(epochs):\n",
        "  \"\"\"Returns an alpha scheduler function for the given configuration.\"\"\"\n",
        "  def scheduler(epoch):\n",
        "    if epoch < 1/4 * epochs:\n",
        "      return 3. * (epoch + 1) / (epochs/4 + 1)\n",
        "    return None\n",
        "  return scheduler\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "VOdc1YWLtWpI"
      },
      "outputs": [],
      "source": [
        "#@title Tensorboard logging callback\n",
        "\n",
        "class LogCallback(tf.keras.callbacks.Callback):\n",
        "  \"\"\"Logs metrics to TensorBoard.\"\"\"\n",
        "\n",
        "  def __init__(self, log_path):\n",
        "    super().__init__()\n",
        "    self.log_path = log_path\n",
        "    self._train_graph = None\n",
        "    self._test_graph = None\n",
        "\n",
        "  def on_train_begin(self, logs=None):\n",
        "    del logs  # unused\n",
        "    if not hasattr(self, \"train_writer\"):\n",
        "      self.train_writer = tf.summary.create_file_writer(\n",
        "          self.log_path + \"/train\")\n",
        "    self.log_variables()\n",
        "\n",
        "  def on_test_begin(self, logs=None):\n",
        "    del logs  # unused\n",
        "    if not hasattr(self, \"test_writer\"):\n",
        "      self.test_writer = tf.summary.create_file_writer(\n",
        "          self.log_path + \"/val\")\n",
        "\n",
        "  def on_test_end(self, logs=None):\n",
        "    # Log test metrics.\n",
        "    self.log_tensorboard(\n",
        "        self.test_writer, {\"metrics/\" + l: v for l, v in logs.items()})\n",
        "    self.test_writer.flush()\n",
        "\n",
        "  def on_epoch_begin(self, epoch, logs=None):\n",
        "    del logs  # unused\n",
        "    self.model.epoch.assign(epoch)\n",
        "\n",
        "  def on_epoch_end(self, epoch, logs=None):\n",
        "    logs = dict(logs)\n",
        "    lr = logs.pop(\"lr\")\n",
        "    alpha = logs.pop(\"alpha\", None)\n",
        "\n",
        "    # Log training metrics.\n",
        "    logs = {l: v for l, v in logs.items() if not l.startswith(\"val_\")}\n",
        "    self.log_tensorboard(\n",
        "        self.train_writer, {\"metrics/\" + l: v for l, v in logs.items()})\n",
        "\n",
        "    # Log learning rate.\n",
        "    logs = {\"learning rate\": lr}\n",
        "    if alpha is not None:\n",
        "      logs[\"alpha\"] = alpha\n",
        "    self.log_tensorboard(self.train_writer, logs)\n",
        "\n",
        "    self.train_writer.flush()\n",
        "\n",
        "  def on_train_batch_begin(self, batch, logs=None):\n",
        "    del logs  # unused\n",
        "    if batch == 0 and not self._train_graph:\n",
        "      with self.train_writer.as_default():\n",
        "        tf.summary.trace_on(graph=True, profiler=False)\n",
        "      self._train_graph = \"tracing\"\n",
        "\n",
        "  def on_train_batch_end(self, batch, logs=None):\n",
        "    del batch, logs  # unused\n",
        "    if self._train_graph == \"tracing\":\n",
        "      with self.train_writer.as_default():\n",
        "        tf.summary.trace_export(\"step\", step=self.model.epoch.value())\n",
        "      self._train_graph = \"traced\"\n",
        "\n",
        "  def on_test_batch_begin(self, batch, logs=None):\n",
        "    del logs  # unused\n",
        "    if batch == 0 and not self._test_graph:\n",
        "      with self.test_writer.as_default():\n",
        "        tf.summary.trace_on(graph=True, profiler=False)\n",
        "      self._test_graph = \"tracing\"\n",
        "\n",
        "  def on_test_batch_end(self, batch, logs=None):\n",
        "    del batch, logs  # unused\n",
        "    if self._test_graph == \"tracing\":\n",
        "      with self.test_writer.as_default():\n",
        "        tf.summary.trace_export(\"step\", step=self.model.epoch.value())\n",
        "      self._test_graph = \"traced\"\n",
        "\n",
        "  def log_tensorboard(self, writer, logs):\n",
        "    \"\"\"Logs the values in `logs` to the summary writer.\"\"\"\n",
        "    with writer.as_default():\n",
        "      for label, value in logs.items():\n",
        "        tf.summary.scalar(label, value, step=self.model.epoch.value())\n",
        "\n",
        "  def log_variables(self):\n",
        "    \"\"\"Logs shape and dtypes of variable collections.\"\"\"\n",
        "    model = self.model\n",
        "    var_format = lambda v: f\"{v.name} {v.dtype} {v.shape}\"\n",
        "    logging.info(\n",
        "        \"TRAINABLE VARIABLES:\\n%s\\n\",\n",
        "        \"\\n\".join(var_format(v) for v in model.trainable_variables))\n",
        " "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JxeKAtQduEe8"
      },
      "source": [
        "# Laplace source"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Y4zrQRkcXDWV"
      },
      "outputs": [],
      "source": [
        "#@title VECVQ\n",
        "\n",
        "work_path = \"/tmp/toy_sources/laplace/vecvq\"\n",
        "\n",
        "epochs = 50\n",
        "steps_per_epoch = 1000\n",
        "batch_size = 1024\n",
        "validation_size = 10000000\n",
        "validation_batch_size = 65536\n",
        "learning_rate = 1e-3\n",
        "\n",
        "codebook_size = 128\n",
        "lmbda = 1.\n",
        "\n",
        "# tf.debugging.enable_check_numerics()\n",
        "\n",
        "source = get_laplace()\n",
        "optimizer = tf.keras.optimizers.Adam()\n",
        "model = get_vecvq_model(\n",
        "    codebook_size=codebook_size, initialize=\"uniform-40\",\n",
        "    source=source, lmbda=lmbda, distortion_loss=\"sse\")\n",
        "model.compile(optimizer=optimizer)\n",
        "\n",
        "# Add an epoch counter for keeping track in checkpoints.\n",
        "model.epoch = tf.Variable(0, trainable=False, dtype=tf.int64)\n",
        "\n",
        "lr_scheduler = get_lr_scheduler(learning_rate, epochs)\n",
        "alpha_scheduler = get_alpha_scheduler(epochs)\n",
        "callback_list = [\n",
        "    tf.keras.callbacks.ModelCheckpoint(\n",
        "        work_path + \"/checkpoints/ckpt-{epoch:04d}\",\n",
        "        save_weights_only=True),\n",
        "    tf.keras.callbacks.BackupAndRestore(\n",
        "        work_path + \"/backup\"),\n",
        "    tf.keras.callbacks.LearningRateScheduler(lr_scheduler),\n",
        "    AlphaScheduler(alpha_scheduler),\n",
        "    LogCallback(work_path),\n",
        "]\n",
        "\n",
        "model.fit(\n",
        "    epochs=epochs,\n",
        "    steps_per_epoch=steps_per_epoch,\n",
        "    batch_size=batch_size,\n",
        "    validation_size=validation_size,\n",
        "    validation_batch_size=validation_batch_size,\n",
        "    verbose=2,\n",
        "    callbacks=tf.keras.callbacks.CallbackList(callback_list, model=model),\n",
        ")\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HmA5gTisumhN"
      },
      "outputs": [],
      "source": [
        "#@title NTC\n",
        "\n",
        "work_path = \"/tmp/toy_sources/laplace/ntc\"\n",
        "\n",
        "epochs = 50\n",
        "steps_per_epoch = 1000\n",
        "batch_size = 1024\n",
        "validation_size = 10000000\n",
        "validation_batch_size = 65536\n",
        "learning_rate = 1e-3\n",
        "\n",
        "latent_dims = 1\n",
        "analysis_filters = [50, 50]\n",
        "analysis_activation = \"softplus\"\n",
        "synthesis_filters = [50, 50]\n",
        "synthesis_activation = \"softplus\"\n",
        "prior_type = \"deep\"\n",
        "dither = (1, 1, 0, 0)\n",
        "soft_round = (1, 0)\n",
        "guess_offset = False\n",
        "lmbda = 1.\n",
        "\n",
        "# tf.debugging.enable_check_numerics()\n",
        "\n",
        "source = get_laplace()\n",
        "optimizer = tf.keras.optimizers.Adam()\n",
        "model = get_ntc_mlp_model(\n",
        "    latent_dims=latent_dims,\n",
        "    analysis_filters=analysis_filters,\n",
        "    analysis_activation=analysis_activation,\n",
        "    synthesis_filters=synthesis_filters,\n",
        "    synthesis_activation=synthesis_activation,\n",
        "    prior_type=prior_type,\n",
        "    dither=dither,\n",
        "    soft_round=soft_round,\n",
        "    guess_offset=guess_offset,\n",
        "    source=source, lmbda=lmbda, distortion_loss=\"sse\")\n",
        "model.compile(optimizer=optimizer)\n",
        "\n",
        "# Add an epoch counter for keeping track in checkpoints.\n",
        "model.epoch = tf.Variable(0, trainable=False, dtype=tf.int64)\n",
        "\n",
        "lr_scheduler = get_lr_scheduler(learning_rate, epochs)\n",
        "alpha_scheduler = get_alpha_scheduler(epochs)\n",
        "callback_list = [\n",
        "    tf.keras.callbacks.ModelCheckpoint(\n",
        "        work_path + \"/checkpoints/ckpt-{epoch:04d}\",\n",
        "        save_weights_only=True),\n",
        "    tf.keras.callbacks.BackupAndRestore(\n",
        "        work_path + \"/backup\"),\n",
        "    tf.keras.callbacks.LearningRateScheduler(lr_scheduler),\n",
        "    AlphaScheduler(alpha_scheduler),\n",
        "    LogCallback(work_path),\n",
        "]\n",
        "\n",
        "model.fit(\n",
        "    epochs=epochs,\n",
        "    steps_per_epoch=steps_per_epoch,\n",
        "    batch_size=batch_size,\n",
        "    validation_size=validation_size,\n",
        "    validation_batch_size=validation_batch_size,\n",
        "    verbose=2,\n",
        "    callbacks=tf.keras.callbacks.CallbackList(callback_list, model=model),\n",
        ")\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "906JtTUFrnvJ"
      },
      "outputs": [],
      "source": [
        "model.plot_quantization([(-5, 5, 1000)])\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "n4WddYVWtqG6"
      },
      "outputs": [],
      "source": [
        "model.plot_transfer([(-5, 5, 1000)])\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cVgn8T6pySgP"
      },
      "source": [
        "# Banana source"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XTb6iFgOySga"
      },
      "outputs": [],
      "source": [
        "#@title VECVQ\n",
        "\n",
        "work_path = \"/tmp/toy_sources/banana/vecvq\"\n",
        "\n",
        "epochs = 100\n",
        "steps_per_epoch = 1000\n",
        "batch_size = 1024\n",
        "validation_size = 10000000\n",
        "validation_batch_size = 65536\n",
        "learning_rate = 1e-3\n",
        "\n",
        "codebook_size = 256\n",
        "lmbda = 1.\n",
        "\n",
        "# tf.debugging.enable_check_numerics()\n",
        "\n",
        "source = get_banana()\n",
        "optimizer = tf.keras.optimizers.Adam()\n",
        "model = get_vecvq_model(\n",
        "    codebook_size=codebook_size, initialize=\"sample\",\n",
        "    source=source, lmbda=lmbda, distortion_loss=\"sse\")\n",
        "model.compile(optimizer=optimizer)\n",
        "\n",
        "# Add an epoch counter for keeping track in checkpoints.\n",
        "model.epoch = tf.Variable(0, trainable=False, dtype=tf.int64)\n",
        "\n",
        "lr_scheduler = get_lr_scheduler(learning_rate, epochs)\n",
        "alpha_scheduler = get_alpha_scheduler(epochs)\n",
        "callback_list = [\n",
        "    tf.keras.callbacks.ModelCheckpoint(\n",
        "        work_path + \"/checkpoints/ckpt-{epoch:04d}\",\n",
        "        save_weights_only=True),\n",
        "    tf.keras.callbacks.BackupAndRestore(\n",
        "        work_path + \"/backup\"),\n",
        "    tf.keras.callbacks.LearningRateScheduler(lr_scheduler),\n",
        "    AlphaScheduler(alpha_scheduler),\n",
        "    LogCallback(work_path),\n",
        "]\n",
        "\n",
        "model.fit(\n",
        "    epochs=epochs,\n",
        "    steps_per_epoch=steps_per_epoch,\n",
        "    batch_size=batch_size,\n",
        "    validation_size=validation_size,\n",
        "    validation_batch_size=validation_batch_size,\n",
        "    verbose=2,\n",
        "    callbacks=tf.keras.callbacks.CallbackList(callback_list, model=model),\n",
        ")\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "htx7oaT7ySga"
      },
      "outputs": [],
      "source": [
        "#@title NTC\n",
        "\n",
        "work_path = \"/tmp/toy_sources/banana/ntc\"\n",
        "\n",
        "epochs = 100\n",
        "steps_per_epoch = 1000\n",
        "batch_size = 1024\n",
        "validation_size = 10000000\n",
        "validation_batch_size = 65536\n",
        "learning_rate = 1e-3\n",
        "\n",
        "latent_dims = 2\n",
        "analysis_filters = [100, 100]\n",
        "analysis_activation = \"softplus\"\n",
        "synthesis_filters = [100, 100]\n",
        "synthesis_activation = \"softplus\"\n",
        "prior_type = \"deep\"\n",
        "dither = (1, 1, 0, 0)\n",
        "soft_round = (1, 0)\n",
        "guess_offset = False\n",
        "lmbda = 1.\n",
        "\n",
        "# tf.debugging.enable_check_numerics()\n",
        "\n",
        "source = get_banana()\n",
        "optimizer = tf.keras.optimizers.Adam()\n",
        "model = get_ntc_mlp_model(\n",
        "    latent_dims=latent_dims,\n",
        "    analysis_filters=analysis_filters,\n",
        "    analysis_activation=analysis_activation,\n",
        "    synthesis_filters=synthesis_filters,\n",
        "    synthesis_activation=synthesis_activation,\n",
        "    prior_type=prior_type,\n",
        "    dither=dither,\n",
        "    soft_round=soft_round,\n",
        "    guess_offset=guess_offset,\n",
        "    source=source, lmbda=lmbda, distortion_loss=\"sse\")\n",
        "model.compile(optimizer=optimizer)\n",
        "\n",
        "# Add an epoch counter for keeping track in checkpoints.\n",
        "model.epoch = tf.Variable(0, trainable=False, dtype=tf.int64)\n",
        "\n",
        "lr_scheduler = get_lr_scheduler(learning_rate, epochs)\n",
        "alpha_scheduler = get_alpha_scheduler(epochs)\n",
        "callback_list = [\n",
        "    tf.keras.callbacks.ModelCheckpoint(\n",
        "        work_path + \"/checkpoints/ckpt-{epoch:04d}\",\n",
        "        save_weights_only=True),\n",
        "    tf.keras.callbacks.BackupAndRestore(\n",
        "        work_path + \"/backup\"),\n",
        "    tf.keras.callbacks.LearningRateScheduler(lr_scheduler),\n",
        "    AlphaScheduler(alpha_scheduler),\n",
        "    LogCallback(work_path),\n",
        "]\n",
        "\n",
        "model.fit(\n",
        "    epochs=epochs,\n",
        "    steps_per_epoch=steps_per_epoch,\n",
        "    batch_size=batch_size,\n",
        "    validation_size=validation_size,\n",
        "    validation_batch_size=validation_batch_size,\n",
        "    verbose=2,\n",
        "    callbacks=tf.keras.callbacks.CallbackList(callback_list, model=model),\n",
        ")\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nE6yr3_lySgb"
      },
      "outputs": [],
      "source": [
        "model.plot_quantization(2 * [(-5, 5, 1000)])\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KmoKt709KFAv"
      },
      "source": [
        "# Sawbridge source"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fEDaxDdIKFAv"
      },
      "outputs": [],
      "source": [
        "#@title VECVQ\n",
        "\n",
        "work_path = \"/tmp/toy_sources/sawbridge/vecvq\"\n",
        "\n",
        "epochs = 200\n",
        "steps_per_epoch = 1000\n",
        "batch_size = 1024\n",
        "validation_size = 10000000\n",
        "validation_batch_size = 4096\n",
        "learning_rate = 1e-3\n",
        "\n",
        "codebook_size = 50\n",
        "lmbda = 1.\n",
        "\n",
        "# tf.debugging.enable_check_numerics()\n",
        "\n",
        "source = get_sawbridge()\n",
        "optimizer = tf.keras.optimizers.Adam()\n",
        "model = get_vecvq_model(\n",
        "    codebook_size=codebook_size, initialize=\"sample-.1\",\n",
        "    source=source, lmbda=lmbda, distortion_loss=\"mse\")\n",
        "model.compile(optimizer=optimizer)\n",
        "\n",
        "# Add an epoch counter for keeping track in checkpoints.\n",
        "model.epoch = tf.Variable(0, trainable=False, dtype=tf.int64)\n",
        "\n",
        "lr_scheduler = get_lr_scheduler(learning_rate, epochs)\n",
        "alpha_scheduler = get_alpha_scheduler(epochs)\n",
        "callback_list = [\n",
        "    tf.keras.callbacks.ModelCheckpoint(\n",
        "        work_path + \"/checkpoints/ckpt-{epoch:04d}\",\n",
        "        save_weights_only=True),\n",
        "    tf.keras.callbacks.BackupAndRestore(\n",
        "        work_path + \"/backup\"),\n",
        "    tf.keras.callbacks.LearningRateScheduler(lr_scheduler),\n",
        "    AlphaScheduler(alpha_scheduler),\n",
        "    LogCallback(work_path),\n",
        "]\n",
        "\n",
        "model.fit(\n",
        "    epochs=epochs,\n",
        "    steps_per_epoch=steps_per_epoch,\n",
        "    batch_size=batch_size,\n",
        "    validation_size=validation_size,\n",
        "    validation_batch_size=validation_batch_size,\n",
        "    verbose=2,\n",
        "    callbacks=tf.keras.callbacks.CallbackList(callback_list, model=model),\n",
        ")\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UiqjCnmzKFAw"
      },
      "outputs": [],
      "source": [
        "#@title NTC\n",
        "\n",
        "work_path = \"/tmp/toy_sources/sawbridge/ntc\"\n",
        "\n",
        "epochs = 200\n",
        "steps_per_epoch = 1000\n",
        "batch_size = 1024\n",
        "validation_size = 10000000\n",
        "validation_batch_size = 4096\n",
        "learning_rate = 1e-3\n",
        "\n",
        "latent_dims = 10\n",
        "analysis_filters = [100, 100]\n",
        "analysis_activation = \"leaky_relu\"\n",
        "synthesis_filters = [100, 100]\n",
        "synthesis_activation = \"leaky_relu\"\n",
        "prior_type = \"deep\"\n",
        "dither = (1, 1, 0, 0)\n",
        "soft_round = (1, 0)\n",
        "guess_offset = False\n",
        "lmbda = 1.\n",
        "\n",
        "# tf.debugging.enable_check_numerics()\n",
        "\n",
        "source = get_sawbridge()\n",
        "optimizer = tf.keras.optimizers.Adam()\n",
        "model = get_ntc_mlp_model(\n",
        "    latent_dims=latent_dims,\n",
        "    analysis_filters=analysis_filters,\n",
        "    analysis_activation=analysis_activation,\n",
        "    synthesis_filters=synthesis_filters,\n",
        "    synthesis_activation=synthesis_activation,\n",
        "    prior_type=prior_type,\n",
        "    dither=dither,\n",
        "    soft_round=soft_round,\n",
        "    guess_offset=guess_offset,\n",
        "    source=source, lmbda=lmbda, distortion_loss=\"mse\")\n",
        "model.compile(optimizer=optimizer)\n",
        "\n",
        "# Add an epoch counter for keeping track in checkpoints.\n",
        "model.epoch = tf.Variable(0, trainable=False, dtype=tf.int64)\n",
        "\n",
        "lr_scheduler = get_lr_scheduler(learning_rate, epochs)\n",
        "alpha_scheduler = get_alpha_scheduler(epochs)\n",
        "callback_list = [\n",
        "    tf.keras.callbacks.ModelCheckpoint(\n",
        "        work_path + \"/checkpoints/ckpt-{epoch:04d}\",\n",
        "        save_weights_only=True),\n",
        "    tf.keras.callbacks.BackupAndRestore(\n",
        "        work_path + \"/backup\"),\n",
        "    tf.keras.callbacks.LearningRateScheduler(lr_scheduler),\n",
        "    AlphaScheduler(alpha_scheduler),\n",
        "    LogCallback(work_path),\n",
        "]\n",
        "\n",
        "model.fit(\n",
        "    epochs=epochs,\n",
        "    steps_per_epoch=steps_per_epoch,\n",
        "    batch_size=batch_size,\n",
        "    validation_size=validation_size,\n",
        "    validation_batch_size=validation_batch_size,\n",
        "    verbose=2,\n",
        "    callbacks=tf.keras.callbacks.CallbackList(callback_list, model=model),\n",
        ")\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WK060O2MKFAw"
      },
      "outputs": [],
      "source": [
        "#@title KLT (dither)\n",
        "\n",
        "work_path = \"/tmp/toy_sources/sawbridge/klt_dither\"\n",
        "\n",
        "epochs = 200\n",
        "steps_per_epoch = 1000\n",
        "batch_size = 1024\n",
        "validation_size = 10000000\n",
        "validation_batch_size = 4096\n",
        "learning_rate = 1e-3\n",
        "\n",
        "num_samples = (1000, 10000)\n",
        "tolerance = 1e-2\n",
        "latent_dims = 50\n",
        "prior_type = \"deep\"\n",
        "dither = (1, 1, 1, 1)\n",
        "soft_round = (0, 0)\n",
        "guess_offset = False\n",
        "lmbda = 1.\n",
        "\n",
        "# tf.debugging.enable_check_numerics()\n",
        "\n",
        "source = get_sawbridge()\n",
        "optimizer = tf.keras.optimizers.Adam()\n",
        "model = get_ltc_klt_model(\n",
        "    num_samples=num_samples,\n",
        "    tolerance=tolerance,\n",
        "    latent_dims=latent_dims,\n",
        "    prior_type=prior_type,\n",
        "    dither=dither,\n",
        "    soft_round=soft_round,\n",
        "    guess_offset=guess_offset,\n",
        "    source=source, lmbda=lmbda, distortion_loss=\"mse\")\n",
        "model.compile(optimizer=optimizer)\n",
        "\n",
        "# Add an epoch counter for keeping track in checkpoints.\n",
        "model.epoch = tf.Variable(0, trainable=False, dtype=tf.int64)\n",
        "\n",
        "lr_scheduler = get_lr_scheduler(learning_rate, epochs)\n",
        "alpha_scheduler = get_alpha_scheduler(epochs)\n",
        "callback_list = [\n",
        "    tf.keras.callbacks.ModelCheckpoint(\n",
        "        work_path + \"/checkpoints/ckpt-{epoch:04d}\",\n",
        "        save_weights_only=True),\n",
        "    tf.keras.callbacks.BackupAndRestore(\n",
        "        work_path + \"/backup\"),\n",
        "    tf.keras.callbacks.LearningRateScheduler(lr_scheduler),\n",
        "    AlphaScheduler(alpha_scheduler),\n",
        "    LogCallback(work_path),\n",
        "]\n",
        "\n",
        "model.fit(\n",
        "    epochs=epochs,\n",
        "    steps_per_epoch=steps_per_epoch,\n",
        "    batch_size=batch_size,\n",
        "    validation_size=validation_size,\n",
        "    validation_batch_size=validation_batch_size,\n",
        "    verbose=2,\n",
        "    callbacks=tf.keras.callbacks.CallbackList(callback_list, model=model),\n",
        ")\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qEcBsva1RIDQ"
      },
      "outputs": [],
      "source": [
        "#@title Daubechies 4-tap\n",
        "\n",
        "work_path = \"/tmp/toy_sources/sawbridge/daub4\"\n",
        "\n",
        "epochs = 200\n",
        "steps_per_epoch = 1000\n",
        "batch_size = 1024\n",
        "validation_size = 10000000\n",
        "validation_batch_size = 4096\n",
        "learning_rate = 1e-3\n",
        "\n",
        "transform = \"db4\"\n",
        "latent_dims = 50\n",
        "prior_type = \"deep\"\n",
        "dither = (1, 1, 1, 1)\n",
        "soft_round = (0, 0)\n",
        "guess_offset = False\n",
        "lmbda = 1.\n",
        "\n",
        "# tf.debugging.enable_check_numerics()\n",
        "\n",
        "source = get_sawbridge()\n",
        "optimizer = tf.keras.optimizers.Adam()\n",
        "model = get_ltc_ortho_model(\n",
        "    transform=transform,\n",
        "    latent_dims=latent_dims,\n",
        "    prior_type=prior_type,\n",
        "    dither=dither,\n",
        "    soft_round=soft_round,\n",
        "    guess_offset=guess_offset,\n",
        "    source=source, lmbda=lmbda, distortion_loss=\"mse\")\n",
        "model.compile(optimizer=optimizer)\n",
        "\n",
        "# Add an epoch counter for keeping track in checkpoints.\n",
        "model.epoch = tf.Variable(0, trainable=False, dtype=tf.int64)\n",
        "\n",
        "lr_scheduler = get_lr_scheduler(learning_rate, epochs)\n",
        "alpha_scheduler = get_alpha_scheduler(epochs)\n",
        "callback_list = [\n",
        "    tf.keras.callbacks.ModelCheckpoint(\n",
        "        work_path + \"/checkpoints/ckpt-{epoch:04d}\",\n",
        "        save_weights_only=True),\n",
        "    tf.keras.callbacks.BackupAndRestore(\n",
        "        work_path + \"/backup\"),\n",
        "    tf.keras.callbacks.LearningRateScheduler(lr_scheduler),\n",
        "    AlphaScheduler(alpha_scheduler),\n",
        "    LogCallback(work_path),\n",
        "]\n",
        "\n",
        "model.fit(\n",
        "    epochs=epochs,\n",
        "    steps_per_epoch=steps_per_epoch,\n",
        "    batch_size=batch_size,\n",
        "    validation_size=validation_size,\n",
        "    validation_batch_size=validation_batch_size,\n",
        "    verbose=2,\n",
        "    callbacks=tf.keras.callbacks.CallbackList(callback_list, model=model),\n",
        ")\n"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [
        "JxeKAtQduEe8",
        "cVgn8T6pySgP",
        "KmoKt709KFAv"
      ],
      "name": "toy_sources.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}