Action Space

The action space is simple and discrete. At each timestep, the agent can choose one of four actions:

• move left,
• move right,
• move down,
• move up.

Formally, the action is represented as a scalar with shape , taking values from the set .

This representation is a general abstraction used throughout Gymnasium. In more complex environments, a single action may encode multiple simultaneous commands. For example, in a game like Super Mario, a player can jump while moving left or right. Such combinations are still treated as a single action by the environment.

Observation Space

The observation returned by the environment represents the agent’s current position on the grid. Since Frozen Lake consists of a finite number of discrete tiles, each tile is assigned a unique integer identifier.

For example, a grid is indexed as:

More generally, the tile index can be computed as:

where:

• where is row index,
• is column index,
• is number of columns in grid.

This discrete state representation makes Frozen Lake particularly well-suited for tabular methods such as Q-learning.

Rewards

The default reward structure is sparse:

• for reaching the goal tile,
• for stepping onto a frozen tile,
• for falling into a hole.

In other words, the agent receives a reward only when it successfully reaches the goal. This sparse reward setting makes the problem deceptively challenging and highlights the importance of exploration in reinforcement learning.

For full details, see the official Frozen Lake environment documentation.

State Space

The state space consists of all discrete tiles on the grid:

Each state uniquely represents the agent’s current position in the lake.

Action Space

At each time step, the agent can choose one of four actions:

These actions correspond to deterministic intentions, even though the actual transition may be stochastic when the lake is slippery.

Transition Dynamics

The transition function defines the probability of moving from state to state after taking action .

• In the non-slippery version of the environment, transitions are deterministic.
• In the slippery version, the intended action may fail, and the agent may move in a perpendicular direction with non-zero probability.

This stochasticity makes Frozen Lake a useful testbed for algorithms that must learn under uncertainty.

Reward Function

The reward function is sparse and simple:

Episodes terminate when the agent reaches the goal or falls into a hole.

import os

# get rid of the audio warnings
os.environ["SDL_AUDIODRIVER"] = "dummy"

from dataclasses import dataclass

import gymnasium as gym
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
from aquarel import load_theme
from gymnasium.envs.toy_text.frozen_lake import generate_random_map
from tqdm.auto import tqdm, trange

%config InlineBackend.figure_formats = {'retina', 'png'}

@dataclass(frozen=True, slots=True)
class Params:
    n_runs: int = 20  # number of runs from scratch
    total_episodes: int = 2_000  # total episodes (# of playthroughs) in the same run
    learning_rate: float = 0.8  # Q-Learning learning rate
    gamma: float = 0.95  # discounting rate
    epsilon: float = 0.1  # probability of exploration vs. exploitation
    proba_frozen: float = 0.9  # probability that a tile is frozen (not a hole)
    is_slippery: bool = False  # enables slipping: 1/3 forward, 1/3 left, 1/3 right
    seed: int = 123  # seed for reproducability

SHOW_PROGRESS: bool = False

Formal Definition

With the intuition in place, we can now describe the Bernoulli multi-armed bandit more formally. A bandit problem is defined as a tuple , where:

• We have machines (or levers) with probabilities .
• At each time step , we take an action on one slot machine and receive a reward .
• is a set of possible actions. The value of an action is expected reward, . If action corresponds to machine , then .
• is the reward function. In a Bernoulli bandit, each pull yields a reward of with probability , and otherwise.

A Bernoulli bandit can be seen as a simplified Marko decision process (MDP) without a state space. The objective is to maximize the total reward . If we knew which action had the biggest reward probability, this would be equivalent to minimizing the regret from not always choosing that optimal action.

Let denote the reward probability of the optimal action :

The expected cumulative regret up to the time is then:

class BernoulliBandit(BaseBandit):
    def __init__(
        self,
        k: int,
        probas: list[float] | npt.NDArray[np.float64] | None = None,
        seed: int | None = None,
    ):
        # sanity check: `probas` needs to be None or of size `n`.
        assert probas is None or len(probas) == k

        self.k = k  # save number of bandits

        self.rng = np.random.default_rng(seed=seed)

        # random probabilities, if they are explicitly defined
        if probas is None:
            probas = self.rng.random(size=self.k)

        # convert to numpy array for easier operations later
        self.probas = np.asarray(probas)

        # in case of Bernoulli MAB, highest probabily is equal to optimal
        self.best_proba = np.max(self.probas)

    def generate_reward(self, i: int) -> int:
        # The player selected the i-th machine.
        return int(self.rng.random() < self.probas[i])

Bandit Strategies

With the bandint problem formally defined, the next question is: how should we choose actions over time? Different strategies encode different assumptions about how exploration should be handled. Broadly, we can distinguish three categories:

• No exploration: always exploit the best-known action (naive and generally poor).
• Random exploration: explore uniformly at random.
• Informed exploration: explore more often when uncertainty is high.

A simple and widely used example of the last category is the -greedy algorithm.

Information State

At the time step , the algorithm maintains:

• empirical action-value estimates ,
• action counts ,

summarizing all past interactions.

The empirical value estimate for action is defined as:

where:

• is the reward received at time step . For a Bernoulli bandit, this is either (success) or (no reward).
• is an indicator function equal to when action was taken at time , and otherwise.
• is the number of times action has been selected:

Policy

The -greedy policy defines a stochastic action-selection rule:

• with probability , the greedy action is selected:
• with probability , an action is selected uniformly at random.

Equivalently, the policy can be written as:

Update Rule

After selecting action and observing reward , the estimate is updated using the new observation.

class EpsilonGreedy(BaseSolver):
    def __init__(self, bandit: BaseBandit, eps: float, init_proba: float = 1.0, seed: int | None = None) -> None:
        """
        eps (float): the probability to explore at each time step.
        init_proba (float): default to be 1.0; optimistic initialization
        """
        super().__init__(bandit)

        assert 0.0 <= eps <= 1.0
        self.eps = eps

        # optimistic initialization
        self.estimates = np.full(self.bandit.k, fill_value=init_proba, dtype=np.float64)

        # define random generator with seed for reproducibility
        self.rng = np.random.default_rng(seed=seed)

    @property
    def estimated_probas(self) -> npt.NDArray[np.float64]:
        return self.estimates

    def run_one_step(self) -> tuple[int, float]:
        # With probability epsilon pick random exploration, or pick the known best lever.

        if self.rng.random() < self.eps:
            # pure random exploration
            i = self.rng.integers(0, self.bandit.k)
        else:
            # greedy selection with random tie-breaking
            candidates = np.flatnonzero(self.estimates == self.estimates.max())
            i = int(self.rng.choice(candidates))

        r = self.bandit.generate_reward(i)
        self.estimates[i] += 1.0 / (self.counts[i] + 1) * (r - self.estimates[i])

        return i, r

Unified Definition

Hoeffding’s Inequality

If we do not want to assign any prior knowledge about the shape of the reward distribution (e.g., Gaussian, exponential), we can rely on Hoeffding’s Inequality. This theorem is applicable on any bounded distribution.

A random variable is said to follow a bounded distribution if all its values lie within a fixed finite interval . In our case, Bernoulli rewards always lie in , so the boundedness assumption is naturally satisfied.

UCB1

From the previous section, we obtained a general form of the confidence bound:

The remaining question is how to choose the threshold probability . Intuitively, as time goes on and we collect more data, we want our confidence bounds to become tighter and failures to become increasingly unlikely. A simple and effective heuristic is to let the failure probability decrease with time.

A common choice is:

which makes the failure probabilities summable over time and enables strong regret guarantees.

Substituting this into the confidence bound gives:

This yields the classic UCB1 algorithm.

At each time step, UCB1 selects the action that maximizes the optimistic estimate of the reward:

Here:

• promotes exploitation,
• the square-root term promotes exploration, shrinking as increases,
• and the term ensures that even rarely chosen actions are revisited occasionally.

class UCB1(BaseSolver):
    def __init__(self, bandit: BaseBandit, init_proba: float = 1.0, seed: int | None = None):
        super().__init__(bandit)
        self.t = 0  # number of time steps
        self.estimates = np.full(shape=self.bandit.k, fill_value=init_proba, dtype=np.float64)
        self.rng = np.random.default_rng(seed)

    @property
    def estimated_probas(self) -> npt.NDArray[np.float64]:
        return self.estimates

    def run_one_step(self) -> tuple[int, float]:
        self.t += 1

        # Pick the best one with consideration of upper confidence bounds.
        ucb = self.estimates + np.sqrt(2 * np.log(self.t) / (1 + self.counts))

        # tie-breaking
        candidates = np.flatnonzero(ucb == ucb.max())
        i = int(self.rng.choice(candidates))

        r = self.bandit.generate_reward(i)

        self.estimates[i] += 1.0 / (self.counts[i] + 1) * (r - self.estimates[i])

        return i, r

Using Distributional Assumptions

For example, suppose we believe that the mean reward of each slot machine follows a Gaussian likelihood, which induces a Gaussian posterior distribution over the mean reward of each action. After observing rewards for a given action , the posterior is characterized by:

• a posterior mean ,
• and a posterior standard deviation

In this case, a natural choice for the upper confidence bound is the upper quantile of the posterior, for instance a 95% confidence bound:

where corresponds to a 95% credible interval for a Gaussian distribution.

The Bayesian UCB action selection rule then becomes:

Interpretation:

• plays the role of exploitation (current best estimate),
• captures uncertainty (how much we still do not know),
• the constant controls how optimistic we are.

Compared to UCB1, where uncertainty depends only on , Bayesian UCB uses the full posterior uncertainty, which often leads to faster learning when the model assumptions are correct.

Key Difference from UCB1

class BayesianUCB(BaseSolver):
    def __init__(
        self, bandit: BaseBandit, c: float = 2, init_a: float = 1, init_b: float = 1, seed: int | None = None
    ) -> None:
        super().__init__(bandit)

        self.c = c
        self._as = np.full(self.bandit.k, fill_value=init_a, dtype=np.float64)
        self._bs = np.full(self.bandit.k, fill_value=init_b, dtype=np.float64)

        self.t = 0
        self.rng = np.random.default_rng(seed)

    @property
    def estimated_probas(self) -> npt.NDArray[np.float64]:
        return self._as / (self._as + self._bs)

    def run_one_step(self) -> tuple[int, float]:
        from scipy.stats import beta

        self.t += 1

        # ensure each arm is tried at least once
        if self.t <= self.bandit.k:
            i = self.t - 1
        else:
            mu = self._as / (self._as + self._bs)  # posterior mean
            sigma = beta.std(self._as, self._bs)  # posterior std Beta(alpha, beta)
            confidence = mu + self.c * sigma

            # tie-breaking
            candidates = np.flatnonzero(confidence == confidence.max())
            i = self.rng.choice(candidates)

        r = self.bandit.generate_reward(i)

        # update Beta posterior for Bernoulli reward
        self._as[i] += r  # successes
        self._bs[i] += 1 - r  # failures

        return i, r

Thompson Sampling for Bernoulli Bandits (Beta-Bernoulli)

In the Bernoulli banding setting, the reward of each action is either or . The conjugate prior for the Bernoulli distribution is the Beta distribution, so we model each action as:

where:

• counts observed successes,
• counts observed failures.

Initially, we typically use a non-informative prior such as:

class ThompsonSampling(BaseSolver):
    def __init__(self, bandit: BaseBandit, init_a: int = 1, init_b: int = 1, seed: int | None = None) -> None:
        super().__init__(bandit)

        self._as = np.full(self.bandit.k, fill_value=init_a, dtype=np.float64)
        self._bs = np.full(self.bandit.k, fill_value=init_b, dtype=np.float64)

        self.rng = np.random.default_rng(seed)

    @property
    def estimated_probas(self) -> npt.NDArray[np.float64]:
        return self._as / (self._as + self._bs)

    def run_one_step(self) -> tuple[int, float]:
        samples = self.rng.beta(self._as, self._bs)

        # tie-breaking
        candidates = np.flatnonzero(samples == samples.max())
        i = int(self.rng.choice(candidates))

        r = self.bandit.generate_reward(i)

        self._as[i] += r
        self._bs[i] += 1 - r

        return i, r

1. Cumulative Regret Over Time (left)

The left subplot shows how much regret each algorithm accumulates over 10,000 time steps. Lower curve indicate better performance.

• Thompson Sampling performs best. Its regret curve rises slowly at first and then flatten, showing that it quickly identifies the optimal arm and almost never leaves it afterward.
• Bayesian UCB is slightly worse but still competitive. Using posterior uncertainty leads to steady improvement without requiring an explicit exploration parameter.
• -greedy suffers more early regret and coverges more slowly, since it explores randomly rather than strategically.
• UCB1 explores aggresively and therefore incurs noticeably higher regret. This is expected in settings where several arms have relatively high reward probabilities, making early optimistic exploration particularly costly.

The qualitative ordering matches classic theoretical results: Thompson Sampling Bayesian UCB -greedy UCB1 for this type of environment.

2. Estimated Reward Probabilities (middle)

The middle subplot show how accurately each method estimates the reward probability of each arm after training. Arms are sorted by their true values, and the dashed line represents perfect estimation.

• Thompson Sampling and Bayesian UCB are concentrated near the diagonal. Their estimates are reasonably accurate even for suboptimal arms.
• -greedy is more scattered. Because it explores randomly and infrequently revisits some arms, several estimates remain biased or underdeveloped.
• UCB1 tends to overestimate some suboptimal arms early on and then underexplore them later. UCB’s deterministic optimism often leads to distinctive estimation bias: Once the bonus term shrinks, there is little incentive to revisit an arm, even if its estimate is wrong.

This subplot highlights key difference: good decision-making does not always require perfectly accurate models, but algorithms that maintain richer uncertainty estimates (Bayesian UCB and Thompson Sampling) tend to form more reliable estimates.

3. Fraction of Pulls per Arm (right)

The right subplot shows how often each algorithm select each action. Here, the behavioral differences are most visible.

• Thompson Sampling plays the best arm almost exclusively, with its bar nearly reaching 100%.
• Bayesian UCB focuses heavily on the best arm but still allocates a small percentage of trials to others due to posterior uncertainty.
• -greedy spreads its attention more broadly. Because exploration is random, even clearly suboptimal arms continue to receive occasional pulls.
• UCB1 revisits several arms during the optimistic exploration phase. Once the bonus term shrinks, it commits strongly to the best arm, but the early exploration leaves a visible footprint.

This suboptimal emphasizes how each strategy allocates exploration effort:

• -greedy: broad, unfocused exploration
• UCB1: early over-exploration, later commitment
• Bayesian UCB: exploration guided by posterior uncertainty
• Thompson: exploration proportional to probabilities of being optimal

Putting It All Together

These three views (regret, estimation accuracy, and action frequencies) provide a comprehensive picture of each algorithm’s strengths and weaknesses:

• Thompson Sampling is consistent and strong: low regret, accurate estimation, and efficient exploration.
• Bayesian UCB offers a pricipled middle ground and performs well when prior structure is appropriate.
• -greedy is simple but wasteful: random exploration leads to both under- and over-exploration.
• UCB1 works as intended, but deterministic optimism causes large early regret when many arms have similar payoffs.

Results shown correspond to a single random seed; while relative performance may vary across runs, the qualitative behavior and average ordering are consistent with theoretical expectations.

Overall, the benchmark illustrates a central message of the exploration-exploitation dilemma: better uncertainty modeling leads to more efficient learning.

Turnkey Solution?

When planning the private cloud, we initially explored turnkey solutions, including NVIDIA DeepOps. Despite its advantages, we built our custom solution for the following reasons. While DeepOps is an excellent turnkey solution with maintained source code on GitHub and offers commercial support, initial setup requires configuration file adjustments, including Ansible scripts for automated (re)configuration of installed Linux distribution. Its complexity discouraged us from further investing our time in tinkering with it.

One of our biggest concerns was the intricate solution that tries to be versatile and “simple”. However, this inevitably leads to hiding functionalities and, in case of issues, jumping around documentation of multiple unrelated internally used tools. Despite proclaimed simplicity, troubleshooting or upgrade problems require manual intervention, where a thorough understanding of Linux, DeepOps, its internal tooling, and their interactions is necessary for system control. Therefore, we decided to start with a minimalist solution and, over time, plan to expand the system to understand the infrastructure’s operation better.

Foundation Infrastructure

In this section, we discuss the foundation infrastructure of our private cloud solution. We’ll go through these building blocks, including the selection of container management tools and resource sharing, which are vital for the operation of the Kubernetes platform.

Operating System: We chose Ubuntu Server based on the Debian Linux distribution for our system. The advantage of widely used Debian-based Linux distributions is the abundance of available knowledge resources and support, making problem-solving more accessible. Among alternatives, like declarative binary reproducible NixOS and RHEL-based distributions, we also considered the Talos distribution specialized for Kubernetes. However, we preferred to stick with Ubuntu Server due to the Talos project’s novelty and associated risks.

Container Management: For container management, we selected ContainerD, also used in the DeepOps solution and officially supported by NVIDIA. It is an open-source tool that implements the CRI interface for communication between the operating system and Kubernetes for efficient and reliable container management.

Data Storage: For data storage, we chose ZFS, which resides on one of the nodes. Although solutions like HDFS, Gluster, Lustre, or Ceph are far more common in the HPC world, they require dedicated infrastructure and tools to offer features offered by ZFS out-of-the-box. Features include checkpoints, data deduplication, compression, a COW (copy-on-write) system to prevent data loss during writing, immunity to silent bit-rot, the ability to use disks as redundancy for mechanical failures, and the use of fast SSD devices as a cache. It also allows easy manual intervention in the event of incidents. However, at the time of writing, ZFS does not stretch across multiple nodes, posing a risk of cluster failure in case of a data-storing node’s malfunction (single point of failure). There is an ongoing effort to implement ZFS’ distributed RAID (dRAID) [src].

To access ZFS storage from Kubernetes, we used the NFS server, which is part of the Linux kernel. We chose NFS because it is one of the few methods that allow multiple containers to bind to the same mounting point (see table).

System Management: For remote management and node configuration, we use Ansible maintained by Red Hat. We selected it due to its prevalence in other significant open-source projects and positive experiences from past projects.

Kubernetes

In Kubernetes, everything operates as a service. These services provide various functionalities that enhance Kubernetes capabilities, such as storage access, CPU and GPU allocation, traffic management, and connecting services within a mesh network.

To support specific functionalities, appropriate services (much like operating system drivers) must be installed. These specialized services, often called “operators” in Kubernetes terminology [src], are essential. They not only deploy and manage functionalities but also respond to issues. Operators enhance Kubernetes by interfacing with standardized and version-controlled APIs.

Put simply, operators are deployed as controller pods (containers) that watch for changes to custom Kubernetes resources and react accordingly. They function as an intermediary layer, implementing application-specific logic that extends Kubernetes beyond its built-in capabilities.

Hardware

When we bought the hardware in early 2022, we chose third-generation AMD EPYC processors. Specifically, we went for the F-series, which has higher base and turbo frequencies — up to 4.0GHz — at the cost of fewer cores. We picked a CPU with the highest available TDP of 280W. We installed server-grade registered error-correcting memory at the highest frequency supported by the processor and populated all eight channels on both processors. Sixteen sticks of RAM in total. Although we considered solutions from Intel, AMD EPYC processors had better price-to-performance ratios.

From the perspective of numerical performance, our significant concern was Intel-optimized libraries, such as Intel MKL, often found in numerical tools. The library has a “bug” that causes non-Intel processors to utilize a slower SSE instead of more advanced AVX vectorization instructions [src]. OpenBLAS is a good alternative but requires some effort to install it. See Anaconda no-mkl package.

We chose two NVMe drives configured in the mirror configuration (RAID-1) for the system drive. We selected six 8TB SSD SATA drives configured in ZFS RAID-Z1 for data storage, which has one drive redundancy. We also chose two A100 GPUs as accelerators.

NVIDIA A100 GPUs come in two form factors: PCI-E and SXM4. The SXM4 proprietary form factor has a higher TDP and high-bandwidth NVLink interconnections between every GPU through NVSwitch hardware. The downside of SXM4 is that it will only support Ampere generation GPUs and require a special motherboard. The PCI-E variant has a lower TDP, and NVLink can only be across two GPUs. However, we decided against vendor lock-in, limiting ourselves to one brand and generation, and went with the PCI-E variant.

We considered the most likely workflow scenarios. We expected most communication to be CPU-to-GPU, with GPUs sliced into several instances via MIG (Multi-Instance GPU). When MIG mode is enabled, each GPU is partitioned into isolated instances that share the physical GPU resources but do not have access to NVLink interconnects. The slicing configuration can be changed at runtime by recreating the MIG instances.

User Experience

After deploying the hardware and software stack, we conducted a month-long live test to stabilize the configuration. During this period, users were informed that we might reboot the system or make significant changes without responsibility for any potential data loss, though we aimed to minimize such occurrences.

We made two key decisions about resource allocation. Users can utilize all available memory and CPU cores. When CPU demand is high, Kubernetes and the operating system manage the scheduling of tasks. In cases of high memory usage, the job consuming the most memory is terminated to protect other running tasks.

Feedback from students and researchers was overwhelmingly positive, highlighting the high speed, numerous cores, ample memory, and dedicated GPU access without interference.

During the testing phase, “testers” identified several issues, which were promptly addressed. These included adding a shared folder with datasets and Jupyter Notebooks, shared package cache, and better persistence of running tasks in JupyterLab.

JupyterHub

JupyterHub has become a crucial component of our research infrastructure, enhancing our workflow significantly. Its smooth integration was largely due to the interface and functionality of JupyterHub, which closely resemble the tools our researchers and students were familiar with. This similarity played a key role in its quick adoption and high user satisfaction.

Upon logging into JupyterHub, users are presented with a list of predefined containers (as shown in Figure 3). Our recent update includes several options:

• A basic minimal working environment.
• A comprehensive data science environment equipped with multiple packages and support for Python, R, and Julia.
• A selection of containers offering GPU instances.

The development environment greets users with a layout similar to modern IDEs, featuring a file explorer on the left and code editor tabs on the right (see Figure 4).

Grafana

For transparent insight into infrastructure availability, the user has read-only access to the Grafana dashboard. Dashboard visualizes computing resource utilization including metrics like total and per-container CPU usage, memory usage per container, GPU utilization, temperature readings, and storage I/O (see Figure 5).

Objective Function

The interaction between the Generator and the Discriminator can be quantified by their objective or loss functions:

• Discriminator’s Objective: For real data , wants near . For generated data , it targets close to . Its objective is:

• Generator’s Objective: aims for to approach , given by:

Both and continuously improve to outperform each other in this game.

Prepare Components

In the upcoming sections, we’ll do the following steps to prepare the development environment:

• Import necessary libraries, primarily PyTorch and Matplotlib.
• Define constants, including project path and seed, for consistency.
• Determine the computational device (e.g., GPU).
• Provide a weight initialization helper function.

from collections.abc import Callable, Sequence
from pathlib import Path
from typing import Any, Final

import joblib
import numpy as np
from matplotlib import pyplot as plt

%config InlineBackend.figure_formats = {'retina', 'png'}

import torch
from torch import Tensor, nn, optim
from torch.utils.data import ConcatDataset, DataLoader, Dataset
from torchinfo import summary
from torchvision import transforms as T
from torchvision.utils import make_grid
from tqdm import tqdm

SEED: Final[int] = 42

PROJECT_PATH = Path(".").resolve()
FIGURE_PATH = PROJECT_PATH / "figures"
DATASET_PATH = Path.home() / "datasets"

# Common constants for all experiments
IMG_DIM: Final[tuple[int, int, int]] = (1, 28, 28)

device = torch.device("cpu")

if torch.cuda.is_available():
    device = torch.device("cuda")

def weights_init(net: nn.Module) -> None:
    for m in net.modules():
        if isinstance(m, nn.Conv2d | nn.ConvTranspose2d):
            nn.init.normal_(m.weight, 0.0, 0.02)
            if m.bias is not None:
                nn.init.constant_(m.bias, 0.0)

        elif isinstance(m, nn.BatchNorm1d | nn.BatchNorm2d):
            nn.init.normal_(m.weight, 1.0, 0.02)
            if m.bias is not None:
                nn.init.constant_(m.bias, 0.0)

        elif isinstance(m, nn.Linear):
            nn.init.normal_(m.weight, 0, 0.02)
            if m.bias is not None:
                nn.init.constant_(m.bias, 0.0)

class Generator(nn.Module):
    def __init__(self, out_dim: Sequence[int], nz: int = 100, ngf: int = 256, alpha: float = 0.2):
        """
        :param out_dim: output image dimension / shape
        :param nz: size of the latent z vector $z$
        :param ngf: size of feature maps (units in the hidden layers) in the generator
        :param alpha: negative slope of leaky ReLU activation
        """
        super().__init__()
        self.out_dim = out_dim
        self.model = nn.Sequential(
            nn.Linear(nz, ngf),
            nn.LeakyReLU(alpha, inplace=True),
            nn.Linear(ngf, 2 * ngf),
            nn.LeakyReLU(alpha, inplace=True),
            nn.Linear(2 * ngf, 4 * ngf),
            nn.LeakyReLU(alpha, inplace=True),
            nn.Linear(4 * ngf, int(np.prod(self.out_dim))),
            nn.Tanh(),
        )

    def forward(self, x: Tensor) -> Tensor:
        x = self.model(x)
        x = torch.reshape(x, (x.size(0), *self.out_dim))
        return x


summary(Generator(out_dim=(1, 28, 28)), input_size=[128, 100])

==========================================================================================
Layer (type:depth-idx)                   Output Shape              Param #
==========================================================================================
Generator                                [128, 1, 28, 28]          --
├─Sequential: 1-1                        [128, 784]                --
│    └─Linear: 2-1                       [128, 256]                25,856
│    └─LeakyReLU: 2-2                    [128, 256]                --
│    └─Linear: 2-3                       [128, 512]                131,584
│    └─LeakyReLU: 2-4                    [128, 512]                --
│    └─Linear: 2-5                       [128, 1024]               525,312
│    └─LeakyReLU: 2-6                    [128, 1024]               --
│    └─Linear: 2-7                       [128, 784]                803,600
│    └─Tanh: 2-8                         [128, 784]                --
==========================================================================================
Total params: 1,486,352
Trainable params: 1,486,352
Non-trainable params: 0
Total mult-adds (Units.MEGABYTES): 190.25
==========================================================================================
Input size (MB): 0.05
Forward/backward pass size (MB): 2.64
Params size (MB): 5.95
Estimated Total Size (MB): 8.63
==========================================================================================

class Discriminator(nn.Module):
    def __init__(self, input_dim: Sequence[int], ndf: int = 128, alpha: float = 0.2):
        super().__init__()
        self.model = nn.Sequential(
            nn.Linear(int(np.prod(input_dim)), 4 * ndf),
            nn.LeakyReLU(alpha, inplace=True),
            nn.Dropout(0.3),
            nn.Linear(4 * ndf, 2 * ndf),
            nn.LeakyReLU(alpha, inplace=True),
            nn.Dropout(0.3),
            nn.Linear(2 * ndf, ndf),
            nn.LeakyReLU(alpha, inplace=True),
            nn.Dropout(0.3),
            nn.Linear(ndf, 1),
            # nn.Sigmoid(),
        )

    def forward(self, x: Tensor) -> Tensor:
        x = torch.reshape(x, (x.size(0), -1))
        return self.model(x)


summary(Discriminator(input_dim=(1, 28, 28)), input_size=[128, 1, 28, 28])

==========================================================================================
Layer (type:depth-idx)                   Output Shape              Param #
==========================================================================================
Discriminator                            [128, 1]                  --
├─Sequential: 1-1                        [128, 1]                  --
│    └─Linear: 2-1                       [128, 512]                401,920
│    └─LeakyReLU: 2-2                    [128, 512]                --
│    └─Dropout: 2-3                      [128, 512]                --
│    └─Linear: 2-4                       [128, 256]                131,328
│    └─LeakyReLU: 2-5                    [128, 256]                --
│    └─Dropout: 2-6                      [128, 256]                --
│    └─Linear: 2-7                       [128, 128]                32,896
│    └─LeakyReLU: 2-8                    [128, 128]                --
│    └─Dropout: 2-9                      [128, 128]                --
│    └─Linear: 2-10                      [128, 1]                  129
==========================================================================================
Total params: 566,273
Trainable params: 566,273
Non-trainable params: 0
Total mult-adds (Units.MEGABYTES): 72.48
==========================================================================================
Input size (MB): 0.40
Forward/backward pass size (MB): 0.92
Params size (MB): 2.27
Estimated Total Size (MB): 3.59
==========================================================================================

Training Loop

The training process is iterative:

• Update Discriminator: With the Generator static, improve the Discriminator’s detection of real vs. fake.
• Update Generator: With a static Discriminator, enhance the Generator’s ability to deceive.

Training continues until the Generator produces almost authentic data. Equilibrium is reached when the Discriminator sees every input as equally likely real or fake, assigning a probability of .

def train_step(
    generator: nn.Module,
    discriminator: nn.Module,
    optim_G: optim.Optimizer,
    optim_D: optim.Optimizer,
    criterion: Callable[[torch.Tensor, torch.Tensor], torch.Tensor],
    real_data: torch.Tensor,
    noise_dim: int,
    device: torch.device,
) -> tuple[float, float]:
    batch_size = real_data.size(0)
    real_data = real_data.to(device, non_blocking=True)

    ### Train Discriminator
    optim_D.zero_grad(set_to_none=True)

    noise = torch.randn(batch_size, noise_dim, device=device)

    output_real = discriminator(real_data)
    real_labels = torch.ones_like(output_real)
    loss_D_real = criterion(output_real, real_labels)

    fake_data = generator(noise)
    output_fake = discriminator(fake_data.detach())
    fake_labels = torch.zeros_like(output_fake)
    loss_D_fake = criterion(output_fake, fake_labels)

    loss_D = (loss_D_real + loss_D_fake) / 2

    loss_D.backward()
    optim_D.step()

    ### Train Generator
    optim_G.zero_grad(set_to_none=True)

    # Freeze D params so autograd does not waste work computing their grads
    for p in discriminator.parameters():
        p.requires_grad_(False)

    fake_data = generator(noise)
    output_fake = discriminator(fake_data)
    target_for_g = torch.ones_like(output_fake)
    loss_G = criterion(output_fake, target_for_g)

    loss_G.backward()
    optim_G.step()

    for p in discriminator.parameters():
        p.requires_grad_(True)

    return loss_G.detach().item(), loss_D.detach().item()

Evaluation

Before evaluation, we configured the learning rate (LR), optimizer’s parameters, batch size, and data loader settings for all experiments. We used the MNIST digits and MNIST fashion datasets for assessment.

OPTIMIZER_LR = 0.0002
L2_NORM = 1e-5
OPTIMIZER_BETAS = (0.5, 0.999)
N_EPOCHS = 100
BATCH_SIZE = 128

g = torch.Generator()
g.manual_seed(SEED)

loader_kwargs = {
    "num_workers": joblib.cpu_count(only_physical_cores=True),
    "pin_memory": True,
    "shuffle": True,
    "batch_size": BATCH_SIZE,
    "prefetch_factor": 2,
    "persistent_workers": True,
    "worker_init_fn": seed_worker,
    "generator": g,
}

def get_mnist_dataset(transform: T.Compose | None = None) -> Dataset:
    from torchvision.datasets import MNIST

    root = str(DATASET_PATH)
    trainset = MNIST(root=root, train=True, download=True, transform=transform)
    testset = MNIST(root=root, train=False, download=True, transform=transform)
    # Combine train and test dataset for more samples.
    dataset = ConcatDataset([trainset, testset])
    return dataset

NOISE_DIM = 100

transform = T.Compose([T.ToTensor(), T.Normalize(0.5, 0.5)])

dataset = get_mnist_dataset(transform=transform)
dataloader = DataLoader(dataset, **loader_kwargs)

# set seed for random generators
set_random_seed(seed=SEED)

# benchmark_noise is used for the animation to show how output evolve on the same vector
benchmark_noise = torch.randn(16 * 16, NOISE_DIM, device=device)

generator = Generator(out_dim=IMG_DIM, nz=NOISE_DIM).to(device)
generator.apply(weights_init)

discriminator = Discriminator(input_dim=IMG_DIM).to(device)
discriminator.apply(weights_init)

optimizer_G = optim.AdamW(
    generator.parameters(),
    lr=OPTIMIZER_LR,
    betas=OPTIMIZER_BETAS,
    weight_decay=L2_NORM,
)

optimizer_D = optim.AdamW(
    discriminator.parameters(),
    lr=OPTIMIZER_LR,
    betas=OPTIMIZER_BETAS,
    weight_decay=L2_NORM,
)

criterion = nn.BCEWithLogitsLoss().to(device)

animation: list[Tensor] = []

g_losses: list[float] = []
d_losses: list[float] = []

for _ in tqdm(range(N_EPOCHS), unit="epochs"):
    generator.train()
    discriminator.train()

    for samples_real, _ in dataloader:
        g_loss, d_loss = train_step(
            generator,
            discriminator,
            optimizer_G,
            optimizer_D,
            criterion,
            samples_real,
            NOISE_DIM,
            device,
        )

        g_losses.append(g_loss)
        d_losses.append(d_loss)

    generator.eval()
    with torch.inference_mode():
        images = generator(benchmark_noise)
        images = images.cpu()

        images = make_grid(images, nrow=16, normalize=True)
        images = (images * 255).clamp(0, 255).to(torch.uint8)

        animation.append(images)

100%|██████████| 100/100 [05:44<00:00,  3.45s/epochs]

NOISE_DIM: int = 100

def get_mnist_fashion_dataset(transform: T.Compose | None = None) -> Dataset:
    from torchvision.datasets import FashionMNIST

    root = str(DATASET_PATH)
    trainset = FashionMNIST(root=root, train=True, download=True, transform=transform)
    testset = FashionMNIST(root=root, train=False, download=True, transform=transform)
    # Combine train and test dataset for more samples.
    dataset = ConcatDataset([trainset, testset])
    return dataset

transform = T.Compose([T.ToTensor(), T.Normalize(0.5, 0.5)])

data = get_mnist_fashion_dataset(transform=transform)
dataloader = DataLoader(data, **loader_kwargs)

# set seed for random generators
set_random_seed(seed=SEED)

# benchmark_noise is used for the animation to show how output evolve on same vector
benchmark_noise = torch.randn(16 * 16, NOISE_DIM, device=device)

generator = Generator(out_dim=IMG_DIM, nz=NOISE_DIM).to(device)
generator.apply(weights_init)

discriminator = Discriminator(input_dim=IMG_DIM).to(device)
discriminator.apply(weights_init)

optimizer_G = optim.AdamW(
    generator.parameters(),
    lr=OPTIMIZER_LR,
    betas=OPTIMIZER_BETAS,
    weight_decay=L2_NORM,
)

optimizer_D = optim.AdamW(
    discriminator.parameters(),
    lr=OPTIMIZER_LR,
    betas=OPTIMIZER_BETAS,
    weight_decay=L2_NORM,
)

criterion = nn.BCEWithLogitsLoss().to(device)

animation = []

g_losses, d_losses = [], []
for _ in tqdm(range(N_EPOCHS), unit="epochs"):
    generator.train()
    discriminator.train()

    for samples_real, _ in dataloader:
        g_loss, d_loss = train_step(
            generator,
            discriminator,
            optimizer_G,
            optimizer_D,
            criterion,
            samples_real,
            NOISE_DIM,
            device,
        )

        g_losses.append(g_loss)
        d_losses.append(d_loss)

    generator.eval()
    with torch.inference_mode():
        images = generator(benchmark_noise)
        images = images.cpu()

        images = make_grid(images, nrow=16, normalize=True)
        images = (images * 255).clamp(0, 255).to(torch.uint8)

        animation.append(images)

100%|██████████| 100/100 [05:47<00:00,  3.47s/epochs]

Setting Up DCGANs

class Generator(nn.Module):
    def __init__(self, nz: int = 100, ngf: int = 32, nc: int = 1):
        """
        :param nz: size of the latent z vector
        :param ngf: size of feature maps in generator
        :param nc: number of channels in the training images.
        """
        super().__init__()
        self.layers = nn.Sequential(
            nn.ConvTranspose2d(nz, 4 * ngf, 4, 1, 0, bias=False),
            nn.BatchNorm2d(4 * ngf),
            nn.ReLU(inplace=True),
            nn.ConvTranspose2d(4 * ngf, 2 * ngf, 3, 2, 1, bias=False),
            nn.BatchNorm2d(2 * ngf),
            nn.ReLU(inplace=True),
            nn.ConvTranspose2d(2 * ngf, ngf, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf),
            nn.ReLU(inplace=True),
            nn.ConvTranspose2d(ngf, nc, 4, 2, 1, bias=False),
            nn.Tanh(),
        )

    def forward(self, x: Tensor) -> Tensor:
        x = torch.reshape(x, (x.size(0), -1, 1, 1))
        return self.layers(x)


summary(Generator(), input_size=(128, 100))

==========================================================================================
Layer (type:depth-idx)                   Output Shape              Param #
==========================================================================================
Generator                                [128, 1, 28, 28]          --
├─Sequential: 1-1                        [128, 1, 28, 28]          --
│    └─ConvTranspose2d: 2-1              [128, 128, 4, 4]          204,800
│    └─BatchNorm2d: 2-2                  [128, 128, 4, 4]          256
│    └─ReLU: 2-3                         [128, 128, 4, 4]          --
│    └─ConvTranspose2d: 2-4              [128, 64, 7, 7]           73,728
│    └─BatchNorm2d: 2-5                  [128, 64, 7, 7]           128
│    └─ReLU: 2-6                         [128, 64, 7, 7]           --
│    └─ConvTranspose2d: 2-7              [128, 32, 14, 14]         32,768
│    └─BatchNorm2d: 2-8                  [128, 32, 14, 14]         64
│    └─ReLU: 2-9                         [128, 32, 14, 14]         --
│    └─ConvTranspose2d: 2-10             [128, 1, 28, 28]          512
│    └─Tanh: 2-11                        [128, 1, 28, 28]          --
==========================================================================================
Total params: 312,256
Trainable params: 312,256
Non-trainable params: 0
Total mult-adds (Units.GIGABYTES): 1.76
==========================================================================================
Input size (MB): 0.05
Forward/backward pass size (MB): 24.26
Params size (MB): 1.25
Estimated Total Size (MB): 25.56
==========================================================================================

class Discriminator(nn.Module):
    def __init__(self, ndf: int = 32, nc: int = 1, alpha: float = 0.2):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Conv2d(nc, ndf, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf),
            nn.LeakyReLU(alpha, inplace=True),
            nn.Conv2d(ndf, 2 * ndf, 4, 2, 1, bias=False),
            nn.BatchNorm2d(2 * ndf),
            nn.LeakyReLU(alpha, inplace=True),
            nn.Conv2d(2 * ndf, 4 * ndf, 3, 2, 1, bias=False),
            nn.BatchNorm2d(4 * ndf),
            nn.LeakyReLU(alpha, inplace=True),
            nn.Conv2d(4 * ndf, 1, 4, 1, 0, bias=False),
            # nn.Sigmoid(),
        )

    def forward(self, x: Tensor) -> Tensor:
        x = self.layers(x)
        x = torch.reshape(x, (x.size(0), -1))
        return x


summary(Discriminator(), input_size=(BATCH_SIZE, 1, 28, 28))

==========================================================================================
Layer (type:depth-idx)                   Output Shape              Param #
==========================================================================================
Discriminator                            [128, 1]                  --
├─Sequential: 1-1                        [128, 1, 1, 1]            --
│    └─Conv2d: 2-1                       [128, 32, 14, 14]         512
│    └─BatchNorm2d: 2-2                  [128, 32, 14, 14]         64
│    └─LeakyReLU: 2-3                    [128, 32, 14, 14]         --
│    └─Conv2d: 2-4                       [128, 64, 7, 7]           32,768
│    └─BatchNorm2d: 2-5                  [128, 64, 7, 7]           128
│    └─LeakyReLU: 2-6                    [128, 64, 7, 7]           --
│    └─Conv2d: 2-7                       [128, 128, 4, 4]          73,728
│    └─BatchNorm2d: 2-8                  [128, 128, 4, 4]          256
│    └─LeakyReLU: 2-9                    [128, 128, 4, 4]          --
│    └─Conv2d: 2-10                      [128, 1, 1, 1]            2,048
==========================================================================================
Total params: 109,504
Trainable params: 109,504
Non-trainable params: 0
Total mult-adds (Units.MEGABYTES): 369.68
==========================================================================================
Input size (MB): 0.40
Forward/backward pass size (MB): 23.46
Params size (MB): 0.44
Estimated Total Size (MB): 24.30
==========================================================================================

Evaluation

NOISE_DIM = 128

transform = T.Compose(
    [
        T.ToTensor(),
        T.Normalize(0.5, 0.5),
    ]
)

data = get_mnist_dataset(transform)
dataloader = DataLoader(data, **loader_kwargs)

# set seed for random generators
set_random_seed()

# benchmark_noise is used for the animation to show how output evolve on same vector
benchmark_noise = torch.randn(16 * 16, NOISE_DIM, device=device)

generator = Generator(nz=NOISE_DIM, ngf=32, nc=IMG_DIM[0]).to(device)
generator.apply(weights_init)

discriminator = Discriminator(ndf=32, nc=IMG_DIM[0]).to(device)
discriminator.apply(weights_init)

optimizer_G = optim.AdamW(
    generator.parameters(),
    lr=OPTIMIZER_LR,
    betas=OPTIMIZER_BETAS,
    weight_decay=L2_NORM,
)

optimizer_D = optim.AdamW(
    discriminator.parameters(),
    lr=OPTIMIZER_LR,
    betas=OPTIMIZER_BETAS,
    weight_decay=L2_NORM,
)

criterion = nn.BCEWithLogitsLoss().to(device)

animation = []

g_losses, d_losses = [], []
for _ in tqdm(range(N_EPOCHS), unit="epochs"):
    generator.train()
    discriminator.train()

    for samples_real, _ in dataloader:
        g_loss, d_loss = train_step(
            generator,
            discriminator,
            optimizer_G,
            optimizer_D,
            criterion,
            samples_real,
            NOISE_DIM,
            device,
        )

        g_losses.append(g_loss)
        d_losses.append(d_loss)

    generator.eval()
    with torch.inference_mode():
        images = generator(benchmark_noise)
        images = images.cpu()

        images = make_grid(images, nrow=16, normalize=True)
        images = (images * 255).clamp(0, 255).to(torch.uint8)

        animation.append(images)

100%|██████████| 100/100 [04:50<00:00,  2.91s/epochs]

NOISE_DIM = 128

transform = T.Compose(
    [
        T.ToTensor(),
        T.Normalize(0.5, 0.5),
    ]
)

data = get_mnist_fashion_dataset(transform)
dataloader = DataLoader(data, **loader_kwargs)

# set seed for random generators
set_random_seed()

# benchmark_noise is used for the animation to show how output evolve on same vector
benchmark_noise = torch.randn(16 * 16, NOISE_DIM, device=device)

generator = Generator(nz=NOISE_DIM, ngf=32, nc=IMG_DIM[0]).to(device)
generator.apply(weights_init)

discriminator = Discriminator(ndf=32, nc=IMG_DIM[0]).to(device)
discriminator.apply(weights_init)

optimizer_G = optim.AdamW(
    generator.parameters(),
    lr=OPTIMIZER_LR,
    betas=OPTIMIZER_BETAS,
    weight_decay=L2_NORM,
)

optimizer_D = optim.AdamW(
    discriminator.parameters(),
    lr=OPTIMIZER_LR,
    betas=OPTIMIZER_BETAS,
    weight_decay=L2_NORM,
)

criterion = nn.BCEWithLogitsLoss().to(device)

animation = []

g_losses, d_losses = [], []
for _ in tqdm(range(N_EPOCHS), unit="epochs"):
    generator.train()
    discriminator.train()

    for samples_real, _ in dataloader:
        g_loss, d_loss = train_step(
            generator,
            discriminator,
            optimizer_G,
            optimizer_D,
            criterion,
            samples_real,
            NOISE_DIM,
            device,
        )

        g_losses.append(g_loss)
        d_losses.append(d_loss)

    generator.eval()
    with torch.inference_mode():
        images = generator(benchmark_noise)
        images = images.cpu()

        images = make_grid(images, nrow=16, normalize=True)
        images = (images * 255).clamp(0, 255).to(torch.uint8)

        animation.append(images)

100%|██████████| 100/100 [04:55<00:00,  2.96s/epochs]

Loss metrics

To effectively implement Neural Style Transfer, we need to quantify how well the generated image matches both the content and style of our source images. This is done using loss metrics. Let’s delve into the specifics of these metrics and how they drive the NST process.

def content_loss_func(target_features: dict[str, Tensor], precomputed_content_features: dict[str, Tensor]) -> Tensor:
    """Calculate content loss metric for give layers."""

    device = next(iter(target_features.values())).device
    content_loss = torch.tensor(0.0, device=device)

    for layer in precomputed_content_features:
        target_feature = target_features[layer]
        content_feature = precomputed_content_features[layer]

        content_loss += F.mse_loss(target_feature, content_feature)

    return content_loss

def gram_matrix(tensor: Tensor) -> Tensor:
    (b, c, h, w) = tensor.size()

    # reshape into (C x (H x W))
    features = tensor.view(b * c, h * w)

    # compute the gram product
    gram = torch.mm(features, features.t())

    return gram

def style_loss_func(
    target_features: dict[str, Tensor], style_features: dict[str, Tensor], precomputed_style_grams: dict[str, Tensor]
) -> Tensor:
    device = next(iter(target_features.values())).device
    style_loss = torch.tensor(0.0, device=device)

    for layer in style_features:
        target_feature = target_features[layer]
        target_gram = gram_matrix(target_feature)

        style_gram = precomputed_style_grams[layer]

        _, c, h, w = target_feature.shape

        weight = STYLE_LAYERS_DEFAULT[layer]
        layer_style_loss = weight * F.mse_loss(target_gram, style_gram) / (c * h * w)
        style_loss += layer_style_loss

    return style_loss

def total_variance_loss_func(target: Tensor) -> Tensor:
    tv_loss = F.l1_loss(target[:, :, :, :-1], target[:, :, :, 1:]) + F.l1_loss(
        target[:, :, :-1, :], target[:, :, 1:, :]
    )

    return tv_loss

Input preparation

Here we specify path to content and style images:

content_path = "./bridge.jpg"
style_path = "./walking-in-the-rain.jpg"

Neural Style Transfer Process

For feature extraction, we’ll leverage VGG19, pre-trained on ImageNet, same as the original authors. Note that we set the model to evaluation mode, ensuring we only use VGG19 to extract features without altering its weights. We also transfer the neural network (NN) to a chosen device, ideally a GPU, for optimal performance.

# We will use a frozen pre-trained VGG neural network for feature extraction.
# In the original paper, authors have used VGG19 (without batch normalization)
model = models.vgg19(weights=models.VGG19_Weights.IMAGENET1K_V1).features


# Authors in the original paper suggested using AvgPool instead of MaxPool
# for more pleasing results. However, changing the pooling also affects
# activation, so the input needs to be scaled (can't find the original source).
class ScaledAvgPool2d(nn.Module):
    def __init__(self, kernel_size, stride, padding=0, scale_factor=2.0):
        super().__init__()
        self.avgpool = torch.nn.AvgPool2d(kernel_size, stride, padding)
        self.scale_factor = scale_factor

    def forward(self, x):
        return self.avgpool(x) * self.scale_factor


# (OPTIONAL) Replace max-pooling layers with custom average pooling layers
# for i, layer in enumerate(model):
#   if isinstance(layer, torch.nn.MaxPool2d):
#       model[i] = ScaledAvgPool2d(kernel_size=2, stride=2, padding=0)

model = model.eval().requires_grad_(False).to(device)

The pretrained VGG model used normalized ImageNet samples for better performance. For effective style transfer, we’ll follow suit to improve feature extraction. Though images will appear altered post-normalization, they are reverted to their original state after the NST process. Next, we’ll transform the content and style images by:

• Loading them from storage.
• Resizing while maintaining aspect ratio.
• Converting to tensors.
• Normalizing using ImageNet weights.

# ImageNet normalization weights per channel
IMAGENET_MEAN = (0.485, 0.456, 0.406)
IMAGENET_STD = (0.229, 0.224, 0.225)

transform = T.Compose(
    [
        T.ToImage(),
        T.Resize(IMG_SIZE),  # Shorter edge of the image will be matched to `IMG_SIZE`
        T.ToDtype(torch.float32, scale=True),
        T.Normalize(IMAGENET_MEAN, IMAGENET_STD),
    ]
)


def load_image(path: str | Path) -> Tensor:
    img = decode_image(str(path))

    # Transform images into tensors
    img: Tensor = transform(img)

    # Add dimension to imitate batch size equal to 1: (C,H,W) -> (B,C,H,W)
    img = img.unsqueeze(0)
    return img

The following code will prepares content , style , and target images. The target image is a clone of the content image and we enable computation of gradients on it.

# The "style" image from which we obtain style
style = load_image(style_path).to(device)

# The "content" image on which we apply style
content = load_image(content_path).to(device)

# The "target" image to store the outcome
target = content.clone().requires_grad_(True).to(device)

The function below retrieves feature maps from designated layers. As shown in Figure 1:

• Content feature map comes from relu5_2.
• Style feature maps are sourced from relu1_1, relu2_1, relu3_1, relu4_1, and relu5_1.

def get_features(image: Tensor, model: nn.Module, layers: Iterable[str] | None = None) -> dict[str, Tensor]:
    if layers is None:
        layers = tuple(STYLE_LAYERS_DEFAULT.keys()) + CONTENT_LAYERS_DEFAULT

    features = {}
    block_num = 1
    conv_num = 0

    x = image

    for layer in model.children():
        x = layer(x)

        if isinstance(layer, nn.Conv2d):
            # produce layer name to find matching convolutions from the paper
            # and store their output for further processing.
            conv_num += 1
            name = f"conv{block_num}_{conv_num}"
            if name in layers:
                features[name] = x

        elif isinstance(layer, nn.MaxPool2d | nn.AvgPool2d | ScaledAvgPool2d):
            # In VGG, each block ends with max/avg pooling layer.
            block_num += 1
            conv_num = 0

        elif isinstance(layer, nn.BatchNorm2d | nn.ReLU):
            pass

        else:
            raise Exception(f"Unknown layer: {layer}")

    return features

Since content and style images never change, we can precompute their feature maps and grams to speed up the NST process.

# Precompute content features, style features, and style gram matrices.
content_features = get_features(content, model, CONTENT_LAYERS_DEFAULT)
style_features = get_features(style, model, STYLE_LAYERS_DEFAULT)

style_grams = {layer: gram_matrix(style_features[layer]) for layer in style_features}

Next, we will use Adam optimizer, where we specify that only target image is considered for optimization.

optimizer = optim.Adam([target], lr=LEARNING_RATE)

The final step of NST is to transfer style using everything we’ve implemented. We extract feature maps, compute total loss, perform steps using gradient descent, and repeat the process N_EPOCHS times. Gradient changes will apply only to the target image.

To notably enhance NST speed, I utilized mixed precision with the unique bfloat16 found in newer hardware. Traditional half-precision float16 doesn’t yield the same results. I’ve tested it. Probably because of the issue with gradient scaling.

pbar = tqdm(range(N_EPOCHS))

for _ in pbar:
    with torch.autocast("cuda", dtype=torch.bfloat16, enabled=AMP_ENABLED):
        target_features = get_features(target, model)

        content_loss = CONTENT_WEIGHT * content_loss_func(target_features, content_features)
        style_loss = STYLE_WEIGHT * style_loss_func(target_features, style_features, style_grams)
        tv_loss = TV_WEIGHT * total_variance_loss_func(target)

        total_loss = content_loss + style_loss + tv_loss

    optimizer.zero_grad()
    total_loss.backward()

    optimizer.step()

    pbar.set_postfix_str(
        f"total_loss={total_loss.item():.2f} "  # noqa: E501
        f"content_loss={content_loss.item():.2f} "
        f"style_loss={style_loss.item():.2f} "
        f"tv_loss={tv_loss.item():.2f}"
    )

100%|██████████| 5000/5000 [01:37<00:00, 51.45it/s, total_loss=43.91 content_loss=8.70 style_loss=29.11 tv_loss=6.11]     

As mentioned before, images need to be denormalized (i.e. reverted back) to correct colors. After that we compare content, style and target images side-by-side.

class InverseNormalize:
    def __init__(self, mean: Sequence[float], std: Sequence[float]) -> None:
        self.mean = torch.as_tensor(mean)
        self.std = torch.as_tensor(std)

    def __call__(self, x_norm: Tensor) -> Tensor:
        # Ensure mean and std have the correct shape
        mean = self.mean.to(x_norm.device).view(-1, 1, 1)
        std = self.std.to(x_norm.device).view(-1, 1, 1)

        # Inverse normalization: x = x_normalized * std + mean
        x = x_norm.mul(std).add(mean)
        return x


class Clip:
    def __init__(self, vmin: float = 0.0, vmax: float = 1.0) -> None:
        self.vmin = vmin
        self.vmax = vmax

    def __call__(self, x: Tensor) -> Tensor:
        return torch.clamp(x, self.vmin, self.vmax)


inv_transform_preview = T.Compose(
    [
        InverseNormalize(IMAGENET_MEAN, IMAGENET_STD),
        T.Resize(IMG_SIZE, antialias=True),
        T.CenterCrop((IMG_SIZE, IMG_SIZE)),
        Clip(),
    ]
)

imgs = [inv_transform_preview(i.detach().squeeze().cpu()) for i in (content, style, target)]

grid = make_grid(imgs)


def show(imgs):
    if not isinstance(imgs, list):
        imgs = [imgs]

    fig, axs = plt.subplots(ncols=len(imgs), figsize=(15, 5), squeeze=False, dpi=92, tight_layout=True, frameon=False)
    for i, img in enumerate(imgs):
        img = img.detach()
        img = VF.to_pil_image(img)
        axs[0, i].imshow(np.asarray(img))
        axs[0, i].set(xticklabels=[], yticklabels=[], xticks=[], yticks=[])


show(grid)

Examples

A few cherry-picked examples of style transfer: