examples/tutorial/GOKU-net_pendulum_tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "668fc98f",
   "metadata": {},
   "source": [
    "# Welcome to the LatentDiffEq.jl tutorial!\n",
    "In this tutorial we will learn how to define a GOKU-net model and train it on pendulum videos. Let's first activate the tutorial project and install all the required packages in their right versions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54db424e",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "using Pkg\n",
    "Pkg.activate(\".\")\n",
    "Pkg.instantiate()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "670c9fbb",
   "metadata": {},
   "source": [
    "The previous step can take a while (~ 10 min) if this is the first time you run it. While waiting, you can watch the JuliaCon 2021 talk which briefly explains the motiviation and main concepts behind the LatentDiffEq.jl package.\n",
    "\n",
    "[![Alt text](https://i.ytimg.com/vi/jhIgs4swrMA/hq720.jpg?sqp=-oaymwEXCNAFEJQDSFryq4qpAwkIARUAAIhCGAE=&rs=AOn4CLDirBL7fLKzeUk4v7OM9LpPUR9CwQ)](https://www.youtube.com/watch?v=jhIgs4swrMA)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c23b095",
   "metadata": {},
   "outputs": [],
   "source": [
    "using LatentDiffEq\n",
    "using ProgressMeter\n",
    "ProgressMeter.ijulia_behavior(:clear)\n",
    "using Random\n",
    "using Statistics\n",
    "using MLDataUtils\n",
    "using Flux.Data: DataLoader\n",
    "using Flux\n",
    "using OrdinaryDiffEq\n",
    "using DiffEqSensitivity\n",
    "using Images\n",
    "using Plots"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb0339fa",
   "metadata": {},
   "source": [
    "## Latent differential equation definition\n",
    "Now that we have all of the packages that we need, let's start to build our model. First we define the latent differential equation that is going to be used in the decoder part of our latent differential equation model. Since we are building the [GOKU-net](https://arxiv.org/abs/2003.10775) model, the latent differential equations are predetermined; in this case, the pendulum equations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3d3b3c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "struct Pendulum{P,S,T,K}\n",
    "\n",
    "    prob::P\n",
    "    solver::S\n",
    "    sensealg::T\n",
    "    kwargs::K\n",
    "\n",
    "    function Pendulum(; kwargs...)\n",
    "        # Parameters and initial conditions only\n",
    "        # used to initialize the ODE problem\n",
    "        u₀ = Float32[1.0, 1.0]\n",
    "        p = Float32[1.]\n",
    "        tspan = (0.f0, 1.f0)\n",
    "\n",
    "        # Define differential equations\n",
    "        function f!(du, u, p, t)\n",
    "                x, y = u\n",
    "                G = 10.0f0\n",
    "                L = p[1]\n",
    "                \n",
    "                du[1] = y\n",
    "                du[2] =  -G/L*sin(x)\n",
    "        end\n",
    "\n",
    "        # Build ODE Problem\n",
    "        prob = ODEProblem(f!, u₀, tspan, p)\n",
    "\n",
    "        # Chose a solver and sensitivity algorithm\n",
    "        solver = Tsit5()\n",
    "        sensalg = ForwardDiffSensitivity()\n",
    "        \n",
    "        P = typeof(prob)\n",
    "        S = typeof(solver)\n",
    "        T = typeof(sensalg)\n",
    "        K = typeof(kwargs)\n",
    "        new{P,S,T,K}(prob, solver, sensalg, kwargs)\n",
    "    end\n",
    "    \n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a8c56287",
   "metadata": {},
   "source": [
    "For more information about defining differential equations you can refer to the [DifferentialEquations.jl tutorials](https://diffeq.sciml.ai/dev/tutorials/ode_example/) and regarding the different sensitivity algorithms choices, you can address the the corresponding [DiffEqFlux.jl documentation](https://diffeqflux.sciml.ai/dev/ControllingAdjoints/).\n",
    "\n",
    "## Model type, latent differential equations and hyperparameters\n",
    "Next we can set the model type and latent differential equation that we are going to use in our model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3cddfd2",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Global model\n",
    "model_type = GOKU_basic()\n",
    "\n",
    "## Latent Differential Equations\n",
    "diffeq = Pendulum() "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce09e486",
   "metadata": {},
   "source": [
    "Plus all of the hyperparameters and training parameters required to train our model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "adc7cd62",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Training params\n",
    "η = 1e-3                        # learning rate\n",
    "decay = 0.001f0                 # decay applied to weights during optimisation\n",
    "batch_size = 64                 # minibatch size\n",
    "seq_len = 50                    # sequence length for training samples\n",
    "epochs = 1500                   # number of epochs for training\n",
    "seed = 3                        # random seed\n",
    "dt = 0.05                       # timestep for ode solve\n",
    "variational = true              # variational or deterministic training\n",
    "\n",
    "## Annealing schedule\n",
    "start_β = 0f0                   # start value\n",
    "end_β = 1f0                     # end value\n",
    "n_cycle = 4                     # number of annealing cycles\n",
    "ratio = 0.9                     # proportion used to increase β (and 1-ratio used to fix β)\n",
    "\n",
    "## Visualization\n",
    "vis_len = 60                    # number of test frames to visualize after each epoch"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d8ecda6",
   "metadata": {},
   "source": [
    "## Data creation: pixel pendulums with different lengths\n",
    "Before beginning training we need to generate the dataset that is going to be fed to the model. To do this we generate a bunch of samples from our defined Pendulum differential equation system and the Luxor.jl package to generate the videos."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f54dc0be",
   "metadata": {},
   "outputs": [],
   "source": [
    "# script for generating pendulum videos\n",
    "include(\"create_data.jl\")\n",
    "\n",
    "latent_data, z₀s, params_data, high_dim_data = generate_dataset(diffeq = diffeq);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "caaceea7",
   "metadata": {},
   "source": [
    "Let's see the dimensions of `high_dim_data`, which contains the pendulum videos."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3d65bbb",
   "metadata": {},
   "outputs": [],
   "source": [
    "@show size(high_dim_data) # different samples\n",
    "@show size(high_dim_data[1]) # time steps\n",
    "@show size(high_dim_data[1][1]); # frames"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "710c43ff",
   "metadata": {},
   "source": [
    "In order to have a better sense of the data that we are working with, let's see some animations of different samples:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec68f437",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "@gif for i ∈ 1:length(high_dim_data[1])\n",
    "    plts = plot.([Gray.(high_dim_data[j][i]) for j in 1:9])\n",
    "    plot(plts..., size = (200,200), axis=nothing)\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "152f5406",
   "metadata": {},
   "source": [
    "Although the true pendulum length (sampled from a uniform distribution between 1 and 2) varies among different samples, the visual length in the images is always the same. The data is created in this way so that the parameters of the differential equations cannot be simply inferred from geometry, but rather are encoded in the dynamics.\n",
    "\n",
    "## Data wrangling \n",
    "Let's do a few preprocessing steps to the videos to make it easier to use the frames to train our GOKU model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3d442ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "# fix random seed for reproducivility\n",
    "Random.seed!(seed) \n",
    "\n",
    "# Stack time for each sample\n",
    "high_dim_data = Flux.stack.(high_dim_data, 3)\n",
    "\n",
    "# Stack all samples\n",
    "high_dim_data = Flux.stack(high_dim_data, 4)\n",
    "h, w, full_seq_len, observations = size(high_dim_data) # 28x28x400x450\n",
    "latent_data = Flux.stack(latent_data, 3)\n",
    "params_data = Flux.stack(params_data, 3);\n",
    "\n",
    "# Vectorize frames\n",
    "high_dim_data = reshape(high_dim_data, :, full_seq_len, observations) # input_dim, time_size, samples\n",
    "high_dim_data = Float32.(high_dim_data)\n",
    "\n",
    "train_set, val_set = Array.(splitobs(high_dim_data, 0.9))\n",
    "train_set_latent, val_set_latent = Array.(splitobs(latent_data, 0.9))\n",
    "train_set_params, val_set_params = Array.(splitobs(params_data, 0.9));\n",
    "\n",
    "loader_train = DataLoader(train_set, batchsize=batch_size, shuffle=true, partial=false)\n",
    "\n",
    "val_set = permutedims(val_set, [1,3,2])\n",
    "t_val = range(0.f0, step=dt, length = size(val_set, 3))\n",
    "input_dim = size(train_set, 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0ac978f0",
   "metadata": {},
   "source": [
    "## Layers definitions\n",
    "In this section we define the model structure. All of the structure framework is summarized in the following diagram with all of the corresponding names for the different layers.\n",
    "\n",
    "![LatentDiffEq.jl_framework](https://user-images.githubusercontent.com/19957518/128906143-41dd1d0a-d081-4261-b413-0327ad5eace2.png)\n",
    "\n",
    "If you wish to use exactly the same architecture as in the original [GOKU-net paper](https://dl.acm.org/doi/abs/10.1145/3450439.3451866) you can conveniently use the `default_layers` function:\n",
    "\n",
    "`encoder_layers, decoder_layers = default_layers(model_type, input_dim, diffeq)`\n",
    "\n",
    "However, let's define each layer manually (using [Flux.jl](https://fluxml.ai/Flux.jl/stable/)) so that you can easily modify them.\n",
    "\n",
    "### Dimensions and hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "862b43bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's set the same dimensions and hyperparameters as in the original GOKU-net implementation\n",
    "hidden_dim_resnet = 200\n",
    "rnn_input_dim = 32\n",
    "rnn_output_dim = 16\n",
    "latent_dim = 16\n",
    "latent_to_diffeq_dim = 200\n",
    "θ_activation = softplus\n",
    "output_activation = σ\n",
    "init = Flux.kaiming_uniform(gain = 1/sqrt(3))\n",
    "\n",
    "z_dim = length(diffeq.prob.u0)\n",
    "θ_dim = length(diffeq.prob.p)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "47af82b3",
   "metadata": {},
   "source": [
    "### Encoder layers\n",
    "First let's go through the Encoder layers one by one.\n",
    "\n",
    "#### Feature Extractor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04745bdb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# resnet\n",
    "l1 = Dense(input_dim, hidden_dim_resnet, relu, init = init)\n",
    "l2 = Dense(hidden_dim_resnet, hidden_dim_resnet, relu, init = init)\n",
    "l3 = Dense(hidden_dim_resnet, hidden_dim_resnet, relu, init = init)\n",
    "l4 = Dense(hidden_dim_resnet, rnn_input_dim, relu, init = init)\n",
    "feature_extractor = Chain(l1,\n",
    "                            SkipConnection(l2, +),\n",
    "                            SkipConnection(l3, +),\n",
    "                            l4)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a2a6ec6",
   "metadata": {},
   "source": [
    "#### Recurrent Pattern Extractor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a720aa0b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# RNN for the initial conditions\n",
    "pe_z₀ = Chain(RNN(rnn_input_dim, rnn_output_dim, relu),\n",
    "                   RNN(rnn_output_dim, rnn_output_dim, relu))\n",
    "\n",
    "# bidirectional LSTM for the parameters\n",
    "pe_θ_forward = Chain(LSTM(rnn_input_dim, rnn_output_dim, init = init),\n",
    "                   LSTM(rnn_output_dim, rnn_output_dim, init = init))\n",
    "\n",
    "pe_θ_backward = Chain(LSTM(rnn_input_dim, rnn_output_dim, init = init),\n",
    "                    LSTM(rnn_output_dim, rnn_output_dim, init = init))\n",
    "\n",
    "pattern_extractor = (pe_z₀, pe_θ_forward, pe_θ_backward)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6883f8b8",
   "metadata": {},
   "source": [
    "#### Latent in"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "939161c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# final fully connected layers before sampling\n",
    "li_μ_z₀ = Dense(rnn_output_dim, latent_dim, init = init)\n",
    "li_logσ²_z₀ = Dense(rnn_output_dim, latent_dim, init = init)\n",
    "\n",
    "li_μ_θ = Dense(rnn_output_dim*2, latent_dim, init = init)\n",
    "li_logσ²_θ = Dense(rnn_output_dim*2, latent_dim, init = init)\n",
    "\n",
    "latent_in = (li_μ_z₀, li_logσ²_z₀, li_μ_θ, li_logσ²_θ)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4aa266a",
   "metadata": {},
   "source": [
    "Once every layer is defined, we bundle them together for later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b08ec5a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "encoder_layers = (feature_extractor, pattern_extractor, latent_in);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e867a715",
   "metadata": {},
   "source": [
    "### Decoder layers\n",
    "Now let's define all of the Decoder layers\n",
    "\n",
    "#### Latent out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7277487",
   "metadata": {},
   "outputs": [],
   "source": [
    "# after sampling in the latent space but before the differential equation layer\n",
    "lo_z₀ = Chain(Dense(latent_dim, latent_to_diffeq_dim, relu, init = init),\n",
    "                    Dense(latent_to_diffeq_dim, z_dim, init = init))\n",
    "\n",
    "lo_θ = Chain(Dense(latent_dim, latent_to_diffeq_dim, relu, init = init),\n",
    "                    Dense(latent_to_diffeq_dim, θ_dim, θ_activation, init = init))\n",
    "\n",
    "latent_out = (lo_z₀, lo_θ)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93d80061",
   "metadata": {},
   "source": [
    "#### Reconstructor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83d6bdd1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# going back to the input space\n",
    "# resnet\n",
    "l1 = Dense(z_dim, hidden_dim_resnet, relu, init = init)\n",
    "l2 = Dense(hidden_dim_resnet, hidden_dim_resnet, relu, init = init)\n",
    "l3 = Dense(hidden_dim_resnet, hidden_dim_resnet, relu, init = init)\n",
    "l4 = Dense(hidden_dim_resnet, input_dim, output_activation, init = init)\n",
    "reconstructor = Chain(l1,\n",
    "                        SkipConnection(l2, +),\n",
    "                        SkipConnection(l3, +),\n",
    "                        l4)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a65dd065",
   "metadata": {},
   "source": [
    "And we bundle them together for later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0ed2324",
   "metadata": {},
   "outputs": [],
   "source": [
    "decoder_layers = (latent_out, diffeq, reconstructor);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7dbacc2",
   "metadata": {},
   "source": [
    "## LatentDiffEqModel creation\n",
    "Here we define the whole model function as a construction of the bundles of layers we defined earlier (encoder_layers and decoder_layers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3891be1e",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = LatentDiffEqModel(model_type, encoder_layers, decoder_layers);\n",
    "ps = Flux.params(model);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42e45b27",
   "metadata": {},
   "source": [
    "## Training tools definitions\n",
    "Here we define all of the parts needed to train our model.\n",
    "\n",
    "### Optimizer\n",
    "We will use [ADAM with weight decay regularization](https://fluxml.ai/Flux.jl/stable/training/optimisers/#Flux.Optimise.ADAMW)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8dfe97db",
   "metadata": {},
   "outputs": [],
   "source": [
    "opt = ADAMW(η,(0.9,0.999), decay)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f0ae752",
   "metadata": {},
   "source": [
    "### KL annealing scheduling\n",
    "Let's set up a [cyclical annealing](https://www.microsoft.com/en-us/research/blog/less-pain-more-gain-a-simple-method-for-vae-training-with-less-of-that-kl-vanishing-agony/):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7805836",
   "metadata": {},
   "outputs": [],
   "source": [
    "annealing_schedule = frange_cycle_linear(epochs, start_β, end_β, n_cycle, ratio)\n",
    "plot(annealing_schedule)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3327cc3",
   "metadata": {},
   "source": [
    "### Loss definition\n",
    "For the loss function, we will use a combination of a reconstruction loss and a KL divergence, just like any VAE-like architecture."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9e799b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "function loss_batch(model, x, t, β, variational)\n",
    "\n",
    "    # Make prediction\n",
    "    X̂, μ, logσ² = model(x, t, variational)\n",
    "    x̂, ẑ, l̂ = X̂\n",
    "\n",
    "    # Compute reconstruction loss\n",
    "    reconstruction_loss = sum(mean((x .- x̂).^2, dims=(2,3)))\n",
    "\n",
    "    # Compute KL loss\n",
    "    kl_loss = vector_kl(μ, logσ²)\n",
    "\n",
    "    return reconstruction_loss + β * kl_loss\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2845cbe0",
   "metadata": {},
   "source": [
    "### Visualization function\n",
    "Let's define a visualization function to randomly pick a sample from the validation set and plot the actual and predicted latent angles and pendulum images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "357fd240",
   "metadata": {},
   "outputs": [],
   "source": [
    "function visualize_val_image(model, val_set, val_set_latent, val_set_params, vis_len, dt, h, w)\n",
    "    \n",
    "    # randomly pick a sample from val_set and a random time interval of length vis_len\n",
    "    j = rand(1:size(val_set, 2))\n",
    "    idxs = rand_time(size(val_set, 3), vis_len)\n",
    "    x = val_set[:, j:j, idxs]\n",
    "    true_latent = val_set_latent[:, idxs, j]\n",
    "    true_params = val_set_params[j]\n",
    "\n",
    "    # create time range for the model diffeq solving\n",
    "    t_val = range(0.f0, step=dt, length=vis_len)\n",
    "\n",
    "    # run model with current parameters on the picked sample\n",
    "    X̂, μ, logσ² = model(x, t_val)\n",
    "    x̂, ẑ, l̂ = X̂\n",
    "    ẑ₀, θ̂ = l̂\n",
    "    θ̂ = θ̂[1]\n",
    "\n",
    "    # plot actual and inferred angles\n",
    "    plt1 = plot(ẑ[1,1,:], legend=false, ylabel=\"inferred angle\", box = :on, color=:indigo, yforeground_color_axis=:indigo, yforeground_color_text=:indigo, yguidefontcolor=:indigo, rightmargin = 2.0Plots.cm)\n",
    "    xlabel!(\"time\")\n",
    "    plt1 = plot!(twinx(), true_latent[1,:], color=:darkorange1, box = :on, xticks=:none, legend=false, ylabel=\"true angle\", yforeground_color_axis=:darkorange1, yforeground_color_text=:darkorange1, yguidefontcolor=:darkorange1)\n",
    "    title!(\"Sample from validation set\")\n",
    "    \n",
    "    # downsample\n",
    "    x = @view x[:, :, 1:6:end]\n",
    "    x̂ = @view x̂[:, :, 1:6:end]\n",
    "\n",
    "    # build frames vectors\n",
    "    to_image(x) = Gray{N0f8}.(reshape(x, h, w))\n",
    "    frames_val = [to_image(xₜ) for xₜ in eachslice(x, dims = 3)]\n",
    "    frames_pred = [to_image(x̂ₜ) for x̂ₜ in eachslice(x̂, dims = 3)]\n",
    "\n",
    "    # plot a mosaic view of the frames\n",
    "    plt2 = mosaicview(frames_val..., frames_pred..., nrow=2, rowmajor=true)\n",
    "    plt2 = plot(plt2, leg = false, ticks = nothing, border = :none)\n",
    "    annotate!((208, -21, (\"True Pendulum Length = $(round(true_params, digits = 2))\", 9, :gray, :right)))\n",
    "    annotate!((208, -11, (\"Inferred Pendulum Length = $(round(θ̂, digits = 2))\", 9, :gray, :right)))\n",
    "    plt = plot(plt1, plt2, layout = @layout([a; b]))\n",
    "    display(plt)\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4139f0d",
   "metadata": {},
   "source": [
    "## Training\n",
    "Finally we arrive at the point of training the model using all of the parts we defined."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29535acf",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "@info \"Start Training of $(typeof(model_type))-net, total $epochs epochs\"\n",
    "for epoch = 1:epochs\n",
    "\n",
    "    # Set annealing factor\n",
    "    β = annealing_schedule[epoch]\n",
    "\n",
    "    # Model evaluation length\n",
    "    t = range(0.f0, step=dt, length=seq_len)\n",
    "\n",
    "    @info \"Epoch $epoch .. (Sequence training length $seq_len)\"\n",
    "    progress = Progress(length(loader_train))\n",
    "\n",
    "    for x in loader_train\n",
    "\n",
    "        # Permute dimesions for having (pixels, batch_size, time)\n",
    "        x = PermutedDimsArray(x, [1,3,2])\n",
    "\n",
    "        # Use only random sequences of length seq_len for the current minibatch\n",
    "        x = time_loader(x, full_seq_len, seq_len)\n",
    "\n",
    "        # Run the model with the current parameters and compute the loss\n",
    "        loss, back = Flux.pullback(ps) do\n",
    "            loss_batch(model, x, t, β, variational)\n",
    "        end\n",
    "        \n",
    "        # Backpropagate and update\n",
    "        grad = back(1f0)\n",
    "        Flux.Optimise.update!(opt, ps, grad)\n",
    "\n",
    "        # Use validation set to get loss and visualisation\n",
    "        val_loss = loss_batch(model, val_set, t_val, β, false)\n",
    "\n",
    "        # Progress meter\n",
    "        next!(progress; showvalues=[(:loss, loss),(:val_loss, val_loss)])\n",
    "    end\n",
    "\n",
    "    visualize_val_image(model, val_set, val_set_latent, val_set_params, vis_len, dt, h, w)\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "555d8a38",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "@webio": {
   "lastCommId": null,
   "lastKernelId": null
  },
  "kernelspec": {
   "display_name": "Julia 1.6.2",
   "language": "julia",
   "name": "julia-1.6"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.6.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}