Neural Style Transfer

Neural Style Transfer

Neural style transfer (Gatys, Ecker, Bethge 2015): combine the content of one image with the style of another. No model training — just iterative optimization of pixel values against a loss defined over a frozen pretrained CNN.

Content + style → synthesized image.

The key insight

In a pretrained ImageNet CNN:

• Deeper layer activations capture content.
• Gram matrices of activations capture style (textures, brush strokes, color palette).

Define a loss matching both; optimize over the synthesized image’s pixels.

Pipeline: forward pass extracts content + style features; backprop into pixels.

Loading content and style

%matplotlib inline
from d2l import jax as d2l
from d2l.nnx_resnet import ResNet50
from flax import nnx
import jax
from jax import numpy as jnp
import optax
import numpy as np
from PIL import Image

d2l.set_figsize()
content_img = Image.open('../img/rainier.jpg')
d2l.plt.imshow(content_img);

style_img = Image.open('../img/autumn-oak.jpg')
d2l.plt.imshow(style_img);

Preprocessing

ImageNet mean/std normalization in, inverse on the way out:

rgb_mean = jnp.array([0.485, 0.456, 0.406])
rgb_std = jnp.array([0.229, 0.224, 0.225])

def preprocess(img, image_shape):
    img = img.resize((image_shape[1], image_shape[0]))  # PIL resize is (w, h)
    img = np.array(img, dtype=np.float32) / 255.0  # (H, W, C)
    img = img.transpose(2, 0, 1)  # (C, H, W)
    img = (img - np.array(rgb_mean).reshape(3, 1, 1)) / np.array(
        rgb_std).reshape(3, 1, 1)
    return jnp.expand_dims(jnp.array(img), axis=0)

def postprocess(img):
    img = np.array(img[0])  # (C, H, W)
    img = np.clip(img.transpose(1, 2, 0) * np.array(rgb_std) +
                  np.array(rgb_mean), 0, 1)
    return img

Pretrained VGG-19 feature extractor

Style is a multi-scale phenomenon — match it across several VGG-19 layers (Conv1_1, 2_1, 3_1, 4_1, 5_1). Content is matched at one deeper layer (Conv4_2):

pretrained_net = ResNet50.from_pretrained()

Downloading bytes:           |  0.00B            
Reconstructing (incomplete total...): |          |  0.00B /  0.00B            
Fetching 1 files:   0%|          | 0/1 [00:00<?, ?it/s]

style_layers, content_layers = [0, 5, 10, 19, 28], [25]

# Layer 0 is the pooled stem; layers 1--4 are the four residual stages.
# Matching all five scales transfers both fine texture and broad structure.
style_layers, content_layers = [0, 1, 2, 3, 4], [3]
net = pretrained_net

Feature extractor (cont.)

def extract_features(X, content_layers, style_layers, model=None):
    """Return selected ResNet features in NCHW layout."""
    model = net if model is None else model
    contents = []
    styles = []
    X = jnp.transpose(X, (0, 2, 3, 1))
    X = nnx.relu(model.stem_bn(model.stem_conv(X)))
    X = nnx.max_pool(X, (3, 3), (2, 2),
                     padding=((1, 1), (1, 1)))
    features = [X]
    for stage in model.stages:
        X = stage(X)
        features.append(X)
    for i, feature in enumerate(features):
        feature = jnp.transpose(feature, (0, 3, 1, 2))
        if i in style_layers:
            styles.append(feature)
        if i in content_layers:
            contents.append(feature)
    return contents, styles

def get_contents(image_shape):
    content_X = preprocess(content_img, image_shape)
    contents_Y, _ = extract_features(content_X, content_layers, style_layers)
    return content_X, contents_Y

def get_styles(image_shape):
    style_X = preprocess(style_img, image_shape)
    _, styles_Y = extract_features(style_X, content_layers, style_layers)
    return style_X, styles_Y

Content loss

Squared error between content and synthesized features at the content layer:

def content_loss(Y_hat, Y):
    return jnp.square(Y_hat - jax.lax.stop_gradient(Y)).mean()

Style loss

Squared error between Gram matrices of features at each style layer. Gram matrix G = F F^\top captures pairwise channel correlations, discarding spatial location:

def gram(X):
    num_channels, n = X.shape[1], d2l.size(X) // X.shape[1]
    X = d2l.reshape(X, (num_channels, n))
    return d2l.matmul(X, d2l.transpose(X)) / (num_channels * n)

def style_loss(Y_hat, gram_Y):
    return jnp.square(gram(Y_hat) - jax.lax.stop_gradient(gram_Y)).mean()

Total variation loss

Penalizes high-frequency noise; keeps the synthesized image smooth:

def tv_loss(Y_hat):
    return 0.5 * (d2l.reduce_mean(
        d2l.abs(Y_hat[:, :, 1:, :] - Y_hat[:, :, :-1, :])) +
                  d2l.reduce_mean(
        d2l.abs(Y_hat[:, :, :, 1:] - Y_hat[:, :, :, :-1])))

Combined loss

\mathcal{L} = \alpha\, \mathcal{L}_\text{content} + \beta\, \mathcal{L}_\text{style} + \gamma\, \mathcal{L}_\text{tv}.

The relative weights determine the visual style — high \beta pushes towards painterly, low \beta keeps photorealism.

content_weight, style_weight, tv_weight = 1, 1e4, 10

def compute_loss(X, contents_Y_hat, styles_Y_hat, contents_Y, styles_Y_gram):
    # Calculate the content, style, and total variance losses respectively
    contents_l = [content_loss(Y_hat, Y) * content_weight for Y_hat, Y in zip(
        contents_Y_hat, contents_Y)]
    styles_l = [style_loss(Y_hat, Y) * style_weight for Y_hat, Y in zip(
        styles_Y_hat, styles_Y_gram)]
    tv_l = tv_loss(X) * tv_weight
    # Add up all the losses
    l = sum(styles_l + contents_l + [tv_l])
    return contents_l, styles_l, tv_l, l

Initializing the synthesized image

Start from the content image (or noise — converges slower but works). The synthesized image is the optimization variable; the network parameters are frozen:

# In JAX, we optimize the synthesized image array directly (no Module needed)

def get_inits(X, lr, styles_Y):
    # Initialize synthesized image to the content image
    gen_img = jnp.array(X, dtype=jnp.float32)
    styles_Y_gram = [gram(Y) for Y in styles_Y]
    return gen_img, styles_Y_gram

Optimization loop

Adam (or LBFGS) optimizes the synthesized image itself. The CNN stays frozen; gradients flow through VGG features back to pixels:

def train(X, contents_Y, styles_Y, lr, num_epochs, lr_decay_epoch):
    X, styles_Y_gram = get_inits(X, lr, styles_Y)
    schedule = optax.exponential_decay(
        lr, transition_steps=lr_decay_epoch, decay_rate=0.8,
        staircase=True)
    optimizer = optax.adam(schedule)
    opt_state = optimizer.init(X)
    animator = d2l.Animator(xlabel='epoch', ylabel='loss',
                            xlim=[10, num_epochs],
                            legend=['content', 'style', 'TV'],
                            ncols=2, figsize=(7, 2.5))

    @nnx.jit
    def train_step(model, X, opt_state):
        def loss_fn(X):
            contents_Y_hat, styles_Y_hat = extract_features(
                X, content_layers, style_layers, model)
            contents_l, styles_l, tv_l, total = compute_loss(
                X, contents_Y_hat, styles_Y_hat, contents_Y,
                styles_Y_gram)
            return total, (jnp.stack(contents_l), jnp.stack(styles_l), tv_l)
        (_, losses), grads = jax.value_and_grad(
            loss_fn, has_aux=True)(X)
        updates, opt_state = optimizer.update(grads, opt_state, X)
        return optax.apply_updates(X, updates), opt_state, losses

    history = []
    for epoch in range(num_epochs):
        X, opt_state, (contents_l, styles_l, tv_l) = train_step(
            net, X, opt_state)
        if (epoch + 1) % 10 == 0:
            animator.axes[1].imshow(postprocess(X))
            animator.add(epoch + 1,
                         [float(jnp.sum(contents_l)),
                          float(jnp.sum(styles_l)),
                          float(tv_l)])
        if (epoch + 1) % 50 == 0:
            history.append((epoch + 1, float(jnp.sum(contents_l)),
                            float(jnp.sum(styles_l)), float(tv_l)))
    for epoch, content_l, style_l, variation_l in history:
        print(f'epoch {epoch}, content {content_l:.3f}, '
              f'style {style_l:.3f}, TV {variation_l:.3f}, '
              f'total {content_l + style_l + variation_l:.3f}')
    return X

Optimization result

After a few hundred iterations, the content layout should remain recognizable while colors and local textures move toward the style image. The three plotted losses are weighted differently, so compare their trends rather than their raw magnitudes:

image_shape = (300, 450)  # PIL Image (h, w)
content_X, contents_Y = get_contents(image_shape)
_, styles_Y = get_styles(image_shape)
output = train(content_X, contents_Y, styles_Y, 0.3, 500, 50)

epoch 50, content 0.875, style 1.601, TV 5.062, total 7.539
epoch 100, content 0.530, style 0.642, TV 3.823, total 4.996
epoch 150, content 0.454, style 0.500, TV 3.160, total 4.114
epoch 200, content 0.421, style 0.431, TV 2.826, total 3.678
epoch 250, content 0.400, style 0.389, TV 2.635, total 3.424
epoch 300, content 0.387, style 0.360, TV 2.513, total 3.260
epoch 350, content 0.377, style 0.339, TV 2.432, total 3.148
epoch 400, content 0.369, style 0.324, TV 2.375, total 3.068
epoch 450, content 0.364, style 0.312, TV 2.333, total 3.008
epoch 500, content 0.359, style 0.303, TV 2.301, total 2.964

Recap

• Style transfer = optimize pixels to minimize a content loss + a Gram-matrix style loss + TV smoothness loss.
• The CNN is frozen; we backprop into the image, not the weights.
• Multi-layer style matching is what gives the recognizable texture-on-content look.
• Modern variants: feedforward style nets (one pass per image), AdaIN, neural style with diffusion models — same idea, faster inference.