Neural Style Transfer

Abstract

Neural Style Transfer (NST) optimizes an image so its high-level content features match one source image while correlations between lower-level features match another. I implement the method described by Gatys et al. in PyTorch and document several practical deviations.

Introduction

Neural Style Transfer (NST) combines the spatial structure of a content image with visual patterns extracted from a style image. It does this by optimizing a target image against features from a frozen convolutional network.

Content is represented by activations from a deeper VGG-19 layer. Style is represented by correlations between activations at several depths, summarized with Gram matrices. The target pixels change until the corresponding content and style losses are small.

Gatys et al. introduced this optimization-based method in “A Neural Algorithm of Artistic Style” (Gatys et al. 2015). I implemented it from scratch to understand the loss functions and feature extraction, with several practical deviations described below. The final gallery contains selected examples rather than a systematic evaluation.

Prerequisites

The implementation uses NumPy, Matplotlib, PyTorch, and Torchvision.

from collections.abc import Iterable, Sequence
from pathlib import Path

import numpy as np
import torch
from matplotlib import pyplot as plt
from torch import Tensor, nn, optim
from torch.nn import functional as F
from torchvision import models
from torchvision.io import decode_image
from torchvision.transforms import functional as VF
from torchvision.transforms import v2 as T
from torchvision.utils import make_grid
from tqdm import tqdm

# Random seed for reproducibility
SEED = 42

# Size of the output image
IMG_SIZE = 512

The code supports CPU and CUDA execution. I ran the stored experiment on an NVIDIA RTX 3090; when PyTorch reports hardware support, the implementation enables bfloat16 autocasting.

AMP_ENABLED = False
COMPILE_ENABLED = True
COMPILE_MODE = "default"

device = torch.device("cpu")

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

    if torch.cuda.is_bf16_supported():
        AMP_ENABLED = True

Implementation

Figure 1: The Neural Style Transfer framework introduced by Gatys et al. distinguishes style and content features from designated layers.

NST differs from a conventional training loop because the network remains frozen and the target image is the optimized parameter. The implementation follows five steps:

• Prepare the content, style, and target images.
• Load a pretrained VGG-19 network and freeze its weights.
• Compute content, style, and total-variation losses.
• Backpropagate through VGG-19 to update only the target image.
• Repeat the update while monitoring the three loss components.

# Gatys et al. used 0.2 for every active style layer. These custom weights
# emphasize the shallow layers, which transferred color more strongly here.
STYLE_LAYERS_DEFAULT = {
    "conv1_1": 0.75,
    "conv2_1": 0.5,
    "conv3_1": 0.2,
    "conv4_1": 0.2,
    "conv5_1": 0.2,
}

# Gatys et al. used conv4_2. I observed no visible difference with conv5_2 in these examples.
CONTENT_LAYERS_DEFAULT = ("conv4_2",)

CONTENT_WEIGHT = 8  # "alpha" in the literature
STYLE_WEIGHT = 70  # "beta" in the literature
TV_WEIGHT = 10  # Optional extension, not part of the original Gatys objective


LEARNING_RATE = 0.004
N_EPOCHS = 5_000

Loss metrics

The optimization requires numerical definitions of content and style similarity. The following losses compare the target with the two source images and regularize its spatial variation.

Content loss metric

Content loss is calculated through Euclidean distance (i.e., mean squared error) between the respective intermediate higher-level feature representation \(F^l\) and \(P^l\) of original input image \(\vec{x}\) and the content image \(\vec{p}\) at layer \(l\).

Hence, a given input image \(\vec{x}\) is encoded in each layer of the CNN by the filter responses to that image. A layer with \(N_l\) distinct filters has \(N_l\) feature maps of size \(M_l\), where \(M_l\) is the height times the width of the feature map. So the response in a layer \(l\) can be stored in a matrix \(F^l \in \mathcal{R}^{N_l \times M_l}\) where \(F_{ij}^{l}\) is the activation of the \(i^{th}\) filter at position \(j\) in layer \(l\).

\[ \mathcal{L}_{content}(\vec{p}, \vec{x}, l) = \frac{1}{2}\sum_{i,j} (F^{l}_{ij} - P^{l}_{ij})^2 \]

Gatys et al. matched content at conv4_2. The implementation below uses that layer but computes mean squared error instead of the summed form above; for a fixed feature shape, this changes the scale rather than the optimum and is absorbed into CONTENT_WEIGHT. I also tested conv5_2 and observed no visible difference in these examples.

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

Style loss

The style loss is more involved than the content loss. We compute it by comparing the Gram matrices of the feature maps from the style image and the generated image.

First, let’s understand the Gram matrix. Given the feature map \(F\) of size \(C \times (H \times W)\), where \(C\) is the number of channels and \(H \times W\) are the spatial dimensions, the Gram matrix \(G\) is of size \(C \times C\) and is computed as

\[ G^l_{ij} = \sum_k F^l_{ik} F^l_{jk} \]

where \(G_{ij}\) is the inner product between vectorized feature maps \(i\) and \(j\). This results in a matrix that captures the correlation between different feature maps and, thus, the style information.

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

Gatys et al. define the style loss between the Gram matrix of the generated image \(G\) and that of the style image \(A\) at layer \(l\) as:

\[ E_l = \frac{1}{4 N^{2}_{l} M^{2}_{l}} \sum_{i,j}(G^l_{ij} - A^l_{ij})^2 \]

Here, \(E_l\) is the style loss for layer \(l\), while \(N_l\) and \(M_l\) are the number of channels and the number of spatial positions in that layer, respectively. \(G_{ij}^l\) and \(A_{ij}^l\) are the Gram matrices of the generated image \(\vec{x}\) and style image \(\vec{a}\).

The total style loss is:

\[ \mathcal{L}_{style}(\vec{a}, \vec{x}) = \sum_{l=0}^{L} w_l E_l \]

For the published results, Gatys et al. assigned \(w_l = 1/5\) to each of conv1_1, conv2_1, conv3_1, conv4_1, and conv5_1.

The implementation below deliberately uses a different normalization:

\[ \widetilde{E}_l = \frac{1}{N_l M_l} \operatorname{MSE}(G^l, A^l) = \frac{1}{N_l^3 M_l} \sum_{i,j}(G^l_{ij} - A^l_{ij})^2. \]

Here, \(\operatorname{MSE}\) averages over the \(N_l^2\) entries of the Gram matrix, which accounts for the additional \(N_l^2\) factor in the equivalent summed form.

This custom normalization gives relatively more influence to shallow, high-resolution layers than the Gatys normalization. Together with the custom layer weights above, it transferred the style image’s color palette more strongly and produced the results I preferred. It should therefore be understood as a practical variation of the original objective, not an exact reproduction of its loss scaling.

def style_loss_func(target_features: dict[str, Tensor], precomputed_style_grams: dict[str, Tensor]) -> Tensor:
    """Calculate the custom style loss for batch size one."""

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

    for layer, weight in STYLE_LAYERS_DEFAULT.items():
        target_feature = target_features[layer]
        target_gram = gram_matrix(target_feature)

        style_gram = precomputed_style_grams[layer]

        _, channels, height, width = target_feature.shape
        spatial_size = height * width

        squared_error = (target_gram - style_gram).square().sum()
        normalization = channels**3 * spatial_size
        layer_style_loss = weight * squared_error / normalization
        style_loss += layer_style_loss

    return style_loss

Total Variation Loss

Total Variation (TV) loss, also known as Total Variation Regularization, is commonly added to Neural Style Transfer to encourage spatial smoothness in the generated image. It was not part of the original Gatys objective and is an optional extension in this implementation. Without it, the output might exhibit noise or oscillations, particularly in regions where the content and style objectives do not offer much guidance.

Given an image \(\vec{x}\) of size \(H \times W \times C\) (height, width, channels), define its horizontal and vertical finite differences as:

\[ \Delta_h \vec{x}_{i,j,c} = x_{i,j+1,c} - x_{i,j,c}, \qquad \Delta_v \vec{x}_{i,j,c} = x_{i+1,j,c} - x_{i,j,c}. \]

Here, \(\Delta_h\) and \(\Delta_v\) denote the horizontal and vertical finite-difference operators, respectively.

The mean anisotropic Total Variation loss is then:

\[ \mathcal{L}_{TV}(\vec{x}) = \operatorname{mean}(|\Delta_h \vec{x}|) + \operatorname{mean}(|\Delta_v \vec{x}|). \]

Here, \(\operatorname{mean}\) averages over all valid spatial positions and color channels.

In simple terms, this loss penalizes abrupt changes in pixel values from one to its neighbors. By minimizing this loss, the generated image becomes smoother, reducing artifacts and unwanted noise. When combined with content and style losses, the TV loss ensures that the resulting image not only captures the content and style of the source images but also looks visually coherent and smooth.

def total_variance_loss_func(target: Tensor) -> Tensor:
    """Calculate mean anisotropic total variation."""

    horizontal_difference = target[:, :, :, 1:] - target[:, :, :, :-1]
    vertical_difference = target[:, :, 1:, :] - target[:, :, :-1, :]

    horizontal_variation = horizontal_difference.abs().mean()
    vertical_variation = vertical_difference.abs().mean()

    tv_loss = horizontal_variation + vertical_variation

    return tv_loss

Total Loss

The objective combines three components:

1. Content loss: matches the target’s deeper activation map to the content image.
2. Style loss: matches Gram matrices from the target and style images.
3. Total variation loss: penalizes abrupt differences between neighboring target pixels.

The resulting total loss is:

\[ \mathcal{L}_{total}(\vec{p},\vec{a},\vec{x}) = \alpha\mathcal{L}_{content}(\vec{p},\vec{x}) + \beta\mathcal{L}_{style}(\vec{a},\vec{x}) + \gamma\mathcal{L}_{TV}(\vec{x}) \]

\(\alpha\), \(\beta\), and \(\gamma\) control the relative influence of content, style, and spatial smoothness. The optimizer changes the target image to minimize their weighted sum.

Input preparation

The experiment uses the following content and style images:

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

Neural Style Transfer Process

As in the original method, feature extraction uses VGG-19 pretrained on ImageNet. Evaluation mode and disabled parameter gradients keep the network fixed while gradients still flow through it to the target image.

Note

Gatys et al. replaced VGG-19’s max pooling with average pooling (Gatys et al. 2015). The stored run here retains the pretrained model’s max-pooling layers. The ScaledAvgPool2d class below records an optional experiment, but its replacement loop is disabled and its scale factor is a heuristic rather than part of the cited method.

# 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-19 weights expect inputs normalized with ImageNet channel statistics. The display code later reverses that normalization. Input preparation therefore consists of:

• 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 loads content image \(\vec{p}\) and style image \(\vec{a}\). It initializes target image \(\vec{x}\) as a clone of the content image and enables gradients only for that target.

# 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 helper records each selected post-ReLU activation under the name of its preceding convolution. The active configuration uses conv4_2 for content and conv1_1, conv2_1, conv3_1, conv4_1, and conv5_1 for style. Figure 1 illustrates the overall flow, but the exact content layer in this implementation is conv4_2.

This naming convention matches CONTENT_LAYERS_DEFAULT and STYLE_LAYERS_DEFAULT above.

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
    current_conv: str | None = None

    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
            current_conv = f"conv{block_num}_{conv_num}"

        elif isinstance(layer, nn.ReLU):
            # VGG uses in-place ReLUs. Capture the activation explicitly
            # instead of relying on mutation of the preceding convolution output.
            if current_conv is not None and current_conv in layers:
                features[current_conv] = 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
            current_conv = None

        elif isinstance(layer, nn.BatchNorm2d):
            pass

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

    return features

Because the content and style images remain fixed, their feature maps and style Gram matrices can be computed once before optimization.

# 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}


@torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=AMP_ENABLED)
def compute_losses(
    target: Tensor,
    content_features: dict[str, Tensor],
    style_grams: dict[str, Tensor],
) -> tuple[Tensor, Tensor, Tensor, Tensor]:
    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_grams)
    tv_loss = TV_WEIGHT * total_variance_loss_func(target)
    total_loss = content_loss + style_loss + tv_loss

    return total_loss, content_loss, style_loss, tv_loss


compute_losses_func = (
    torch.compile(compute_losses, fullgraph=True, mode=COMPILE_MODE) if COMPILE_ENABLED else compute_losses
)

The Adam optimizer receives only target image \(\vec{x}\), so it cannot update VGG-19 or either source image.

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

Each of the N_EPOCHS iterations extracts target features, computes the weighted losses, and updates the target image. Gradients pass through VGG-19, but only the target is registered with the optimizer.

On supported CUDA hardware, the loss function runs under bfloat16 autocasting. I did not benchmark its speedup here. A float16 trial produced a visibly different result, but this experiment did not isolate the cause, so the implementation enables only bfloat16.

pbar = tqdm(range(N_EPOCHS))

for _ in pbar:
    total_loss, content_loss, style_loss, tv_loss = compute_losses_func(target, content_features, style_grams)

    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:16<00:00, 65.35it/s, total_loss=101.05 content_loss=58.08 style_loss=36.85 tv_loss=6.12]     

Before display, inverse normalization restores the image channel ranges. The content, style, and target images can then be compared 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)


DPI = 200  # 2x the ~800px content column, so figures stay crisp on HiDPI screens


def show(imgs, save_to: str | Path) -> None:
    if not isinstance(imgs, list):
        imgs = [imgs]

    fig, axs = plt.subplots(
        ncols=len(imgs), figsize=(15, 5), squeeze=False, dpi=92, constrained_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=[])

    fig.savefig(save_to, dpi=DPI, pil_kwargs={"quality": 90, "alpha_quality": 100, "method": 6})
    plt.close(fig)


show(grid, "./figures/style-transfer-result.webp")

Neural style transfer result: content image (left), style image (center), and optimized target image (right).

Conclusions

Implementing NST clarified its central idea for me: a frozen classifier can define a useful image objective even though it was not trained to generate images. Deeper activations constrain spatial content, Gram matrices constrain feature correlations, and gradients update the image rather than the network.

The result depends strongly on the selected layers and on the content, style, and total-variation weights. The gallery shows selected outcomes from this implementation, so it demonstrates the mechanism but does not compare configurations systematically.

Modern text-to-image systems solve a broader problem than transferring style between two supplied images, so they are not direct replacements for this method. NST remains a compact example of optimizing an input through a frozen feature extractor.

Acknowledgements

Helpful articles and code repositories while writing my implementation:

• Gregor Koehler et al. gkoehler/pytorch-neural-style-transfer (best resource in my opinion)
• Ritul’s Medium article (good resource)
• Pragati Baheti blog visually present style extraction
• Aleksa Gordic (gordicaleksa/pytorch-neural-style-transfer)
• ProGamerGov/neural-style-pt
• Katherine Crowson (rowsonkb/style-transfer-pytorch)
• Derrick Mwiti’s Medium article
• Aman Kumar Mallik’s Medium article

The content and style images used throughout this post and in the Appendix gallery come from the following artists and sources. Most were collected from my earlier gcerar/pytorch-neural-style-transfer repository, which credits the same sources:

• “Gray Bridge and Trees”, Martin Damboldt (the content image)
• “The Persistence of Memory”, Salvador Dali
• “The Scream”, Edvard Munch
• “Udnie”, Francis Picabia
• “Edtaonisl”, Francis Picabia
• “Hand with Reflecting Sphere”, M. C. Escher
• “Mysterious Rain Princess”, Leonid Afremov
• “Walking in the Rain”, Leonid Afremov
• “La Muse”, Pablo Picasso
• “Seated Nude”, Pablo Picasso
• “Composition VII”, Wassily Kandinsky
• “Under the Wave off Kanagawa”, Katsushika Hokusai
• “The Starry Night”, Vincent van Gogh
• “The Night Cafe”, Vincent van Gogh
• “Mondrian World Map”, Michael Tompsett
• “June Tree”, Natasha (Wescoat) Bouchillion
• “Doomguy”, id Software
• Colorful whirlpool, fractal pattern, bamboo forest, feathers, flowers, the “lady” portrait, and “Paul of Tarsus”: artist unknown

Appendix

Source images

The content photo and every style image credited above:

Bamboo forest, artist unknown

Colorful whirlpool, artist unknown

Mondrian World Map, Michael Tompsett

Composition VII, Wassily Kandinsky

Doomguy, id Software

Gray Bridge and Trees, Martin Damboldt (content image)

The Starry Night, Vincent van Gogh

The Persistence of Memory, Salvador Dali

Fractal pattern, artist unknown

Walking in the Rain, Leonid Afremov

Under the Wave off Kanagawa, Katsushika Hokusai

The Night Cafe, Vincent van Gogh

La Muse, Pablo Picasso

Lady, artist unknown

Edtaonisl, Francis Picabia

June Tree, Natasha (Wescoat) Bouchillion

Flowers, artist unknown

Udnie, Francis Picabia

Mysterious Rain Princess, Leonid Afremov

Feathers, artist unknown

The Scream, Edvard Munch

Paul of Tarsus, artist unknown

Seated Nude, Pablo Picasso

Hand with Reflecting Sphere, M. C. Escher

Examples

Every style image credited above, applied to the bridge content image:

References

Gatys, Leon A, Alexander S Ecker, and Matthias Bethge. 2015. “A Neural Algorithm of Artistic Style.” arXiv Preprint arXiv:1508.06576.