19.11  Fully Convolutional Networks

As discussed in Section 19.9, semantic segmentation classifies images in pixel level. A fully convolutional network (FCN) uses a convolutional neural network to transform image pixels to pixel classes (Long et al. 2015). Unlike the CNNs that we encountered earlier for image classification or object detection, a fully convolutional network transforms the height and width of intermediate feature maps back to those of the input image: this is achieved by the transposed convolutional layer introduced in Section 19.10. As a result, the classification output and the input image have a one-to-one correspondence in pixel level: the channel dimension at any output pixel holds the classification results for the input pixel at the same spatial position.

%matplotlib inline
from d2l import torch as d2l
import torch
import torchvision
from torch import nn
from torch.nn import functional as F
%matplotlib inline
from d2l import tensorflow as d2l
import tensorflow as tf
import keras
import numpy as np
from PIL import Image
%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
%matplotlib inline
from d2l import mxnet as d2l
from mxnet import gluon, image, init, np, npx
from mxnet.gluon import nn

npx.set_np()

19.11.1 The Model

Here we describe the basic design of the fully convolutional network model. As shown in Figure 19.11.1, this model first uses a CNN to extract image features, then transforms the number of channels into the number of classes via a \(1\times 1\) convolutional layer, and finally transforms the height and width of the feature maps to those of the input image via the transposed convolution introduced in Section 19.10. As a result, the model output has the same height and width as the input image, where the output channel contains the predicted classes for the input pixel at the same spatial position.

Figure 19.11.1: Fully convolutional network.

Below, we use a ResNet model pretrained on the ImageNet dataset to extract image features and denote the model instance as pretrained_net. The last few layers of this model include a global average pooling layer and a fully connected layer: they are not needed in the fully convolutional network.

pretrained_net = torchvision.models.resnet18(
    weights=torchvision.models.ResNet18_Weights.DEFAULT)
list(pretrained_net.children())[-3:]
[Sequential(
   (0): BasicBlock(
     (conv1): Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
     (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
     (relu): ReLU(inplace=True)
     (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
     (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
     (downsample): Sequential(
       (0): Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)
       (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
     )
   )
   (1): BasicBlock(
     (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
     (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
     (relu): ReLU(inplace=True)
     (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
     (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
   )
 ),
 AdaptiveAvgPool2d(output_size=(1, 1)),
 Linear(in_features=512, out_features=1000, bias=True)]
# Note: keras.applications does not bundle a ResNet-18; it only ships
# ResNet-50/101/152. To match the PT/MX tabs (and the prose), we build a
# ResNet-18 from scratch as a Functional model. Conceptually treat its
# weights as if they had been initialized from ImageNet pretraining; in
# practice you would port pretrained weights from PyTorch.
def _resnet_block(x, num_channels, strides=1, use_1x1conv=False):
    y = keras.layers.Conv2D(num_channels, 3, strides=strides,
                            padding='same', use_bias=False)(x)
    y = keras.layers.BatchNormalization()(y)
    y = keras.layers.ReLU()(y)
    y = keras.layers.Conv2D(num_channels, 3, strides=1,
                            padding='same', use_bias=False)(y)
    y = keras.layers.BatchNormalization()(y)
    if use_1x1conv:
        x = keras.layers.Conv2D(num_channels, 1, strides=strides,
                                use_bias=False)(x)
        x = keras.layers.BatchNormalization()(x)
    return keras.layers.ReLU()(y + x)

inputs = keras.Input(shape=(None, None, 3))
x = keras.layers.Conv2D(64, 7, strides=2, padding='same',
                        use_bias=False)(inputs)
x = keras.layers.BatchNormalization()(x)
x = keras.layers.ReLU()(x)
x = keras.layers.MaxPool2D(pool_size=3, strides=2, padding='same')(x)
x = _resnet_block(x, 64)
x = _resnet_block(x, 64)
x = _resnet_block(x, 128, strides=2, use_1x1conv=True)
x = _resnet_block(x, 128)
x = _resnet_block(x, 256, strides=2, use_1x1conv=True)
x = _resnet_block(x, 256)
x = _resnet_block(x, 512, strides=2, use_1x1conv=True)
features = _resnet_block(x, 512)
# Mirror the structure of torchvision.models.resnet18: include the global
# avg pool + dense head so we can slice them off below.
pooled = keras.layers.GlobalAveragePooling2D()(features)
logits = keras.layers.Dense(1000)(pooled)
pretrained_net = keras.Model(inputs=inputs, outputs=logits)
# Show the last few layers (matches the spirit of the PT/MX displays)
pretrained_net.layers[-3:]
[<ReLU name=re_lu_16, built=True>,
 <GlobalAveragePooling2D name=global_average_pooling2d, built=True>,
 <Dense name=dense, built=True>]
pretrained_net = ResNet50.from_pretrained()
dummy = jnp.ones((1, 320, 480, 3))
print('Feature extractor output shape:',
      pretrained_net.feature_map(dummy).shape)
Downloading bytes:           |  0.00B            
Reconstructing (incomplete total...): |          |  0.00B /  0.00B            
Fetching 1 files:   0%|          | 0/1 [00:00<?, ?it/s]
Feature extractor output shape: (1, 10, 15, 2048)
pretrained_net = gluon.model_zoo.vision.resnet18_v2(pretrained=True)
pretrained_net.features[-3:], pretrained_net.output
(HybridSequential(
   (0): Activation(relu)
   (1): GlobalAvgPool2D(size=(1, 1), stride=(1, 1), padding=(0, 0), ceil_mode=True, global_pool=True, pool_type=avg, layout=NCHW)
   (2): Flatten
 ),
 Dense(512 -> 1000, linear))

Next, we create the fully convolutional network instance net. It copies all the pretrained convolutional layers in the ResNet except for the final global average pooling layer and the fully connected layer that are closest to the output.

net = nn.Sequential(*list(pretrained_net.children())[:-2])
# Build the FCN feature extractor: all layers up to (but not including)
# the global average pooling and dense head — i.e., the full conv body.
# The last conv-block output (`features`) is the 1/32-resolution feature
# map; we use it as the new model output, dropping GAP + Dense.
net = keras.Model(inputs=pretrained_net.input, outputs=features)
# `feature_map` omits the global average pooling and classification head.
net = nn.HybridSequential()
for layer in pretrained_net.features[:-2]:
    net.add(layer)

Given an input with height and width of 320 and 480 respectively, the forward propagation of net reduces the input height and width to 1/32 of the original, namely 10 and 15.

X = torch.rand(size=(1, 3, 320, 480))
net(X).shape
torch.Size([1, 512, 10, 15])
X = tf.random.uniform(shape=(1, 320, 480, 3))
net(X).shape
TensorShape([1, 10, 15, 512])
X = jnp.ones((1, 320, 480, 3))
pretrained_net.feature_map(X).shape
(1, 10, 15, 2048)
X = np.random.uniform(size=(1, 3, 320, 480))
net(X).shape
(1, 512, 10, 15)

Next, we use a \(1\times 1\) convolutional layer to transform the number of output channels into the number of classes (21) of the Pascal VOC2012 dataset. Finally, we need to increase the height and width of the feature maps by 32 times to change them back to the height and width of the input image. Recall how to calculate the output shape of a convolutional layer in Section 6.3. Since \((320-64+16\times2+32)/32=10\) and \((480-64+16\times2+32)/32=15\), we construct a transposed convolutional layer with stride of \(32\), setting the height and width of the kernel to \(64\), the padding to \(16\). In general, we can see that for stride \(s\), padding \(s/2\) (assuming \(s/2\) is an integer), and the height and width of the kernel \(2s\), the transposed convolution will increase the height and width of the input by \(s\) times.

num_classes = 21
net.add_module('final_conv', nn.Conv2d(512, num_classes, kernel_size=1))
net.add_module('transpose_conv', nn.ConvTranspose2d(num_classes, num_classes,
                                    kernel_size=64, padding=16, stride=32))
num_classes = 21
# 1x1 conv to reduce channels to num_classes
final_conv = keras.layers.Conv2D(num_classes, kernel_size=1,
                                 kernel_initializer='glorot_uniform')
# Transposed conv: stride=32, kernel=64, padding='same' upsamples 32x
# (for input height/width divisible by 32, output equals input spatial size)
transpose_conv = keras.layers.Conv2DTranspose(
    num_classes, kernel_size=64, strides=32, padding='same', use_bias=False)

inputs = net.input
x = net.output
x = final_conv(x)
x = transpose_conv(x)
fcn_net = keras.Model(inputs=inputs, outputs=x)
print('FCN output shape:', fcn_net(tf.random.uniform((1, 320, 480, 3))).shape)
FCN output shape: (1, 320, 480, 21)
num_classes = 21

class FCN(nnx.Module):
    """Fully Convolutional Network for semantic segmentation."""
    def __init__(self, backbone, num_classes, *, rngs):
        self.backbone = backbone
        self.classifier = nnx.Conv(2048, num_classes, (1, 1), rngs=rngs)
        self.upsample = nnx.ConvTranspose(
            num_classes, num_classes, (64, 64), strides=32, padding='SAME',
            use_bias=False, rngs=rngs)

    def __call__(self, X):
        X = self.backbone.feature_map(X)
        return self.upsample(self.classifier(X))

net = FCN(pretrained_net, num_classes, rngs=nnx.Rngs(0))
print('FCN output shape:',
      net(jnp.ones((1, 320, 480, 3))).shape)
FCN output shape: (1, 320, 480, 21)
num_classes = 21
net.add(nn.Conv2D(num_classes, kernel_size=1),
        nn.Conv2DTranspose(
            num_classes, kernel_size=64, padding=16, strides=32))

19.11.2 Initializing Transposed Convolutional Layers

We already know that transposed convolutional layers can increase the height and width of feature maps. In image processing, we may need to scale up an image, i.e., upsampling. Bilinear interpolation is one of the commonly used upsampling techniques. It is also often used for initializing transposed convolutional layers.

To explain bilinear interpolation, say that given an input image we want to calculate each pixel of the upsampled output image. In order to calculate the pixel of the output image at coordinate \((x, y)\), first map \((x, y)\) to coordinate \((x', y')\) on the input image, for example, according to the ratio of the input size to the output size. Note that the mapped \(x'\) and \(y'\) are real numbers. Then, find the four pixels closest to coordinate \((x', y')\) on the input image. Finally, the pixel of the output image at coordinate \((x, y)\) is calculated based on these four closest pixels on the input image and their relative distance from \((x', y')\).

Upsampling of bilinear interpolation can be implemented by the transposed convolutional layer with the kernel constructed by the following bilinear_kernel function. Due to space limitations, we only provide the implementation of the bilinear_kernel function below without discussions on its algorithm design.

def bilinear_kernel(in_channels, out_channels, kernel_size):
    factor = (kernel_size + 1) // 2
    if kernel_size % 2 == 1:
        center = factor - 1
    else:
        center = factor - 0.5
    og = (torch.arange(kernel_size).reshape(-1, 1),
          torch.arange(kernel_size).reshape(1, -1))
    filt = (1 - torch.abs(og[0] - center) / factor) * \
           (1 - torch.abs(og[1] - center) / factor)
    weight = torch.zeros((in_channels, out_channels,
                          kernel_size, kernel_size))
    weight[range(in_channels), range(out_channels), :, :] = filt
    return weight
def bilinear_kernel(in_channels, out_channels, kernel_size):
    factor = (kernel_size + 1) // 2
    if kernel_size % 2 == 1:
        center = factor - 1
    else:
        center = factor - 0.5
    og = (np.arange(kernel_size).reshape(-1, 1),
          np.arange(kernel_size).reshape(1, -1))
    filt = (1 - np.abs(og[0] - center) / factor) * \
           (1 - np.abs(og[1] - center) / factor)
    # Keras Conv2DTranspose uses HWIO kernel format (height, width, out, in)
    weight = np.zeros((kernel_size, kernel_size, out_channels, in_channels),
                      dtype=np.float32)
    for i in range(min(in_channels, out_channels)):
        weight[:, :, i, i] = filt
    return weight
def bilinear_kernel(in_channels, out_channels, kernel_size):
    factor = (kernel_size + 1) // 2
    if kernel_size % 2 == 1:
        center = factor - 1
    else:
        center = factor - 0.5
    og = (np.arange(kernel_size).reshape(-1, 1),
          np.arange(kernel_size).reshape(1, -1))
    filt = (1 - np.abs(og[0] - center) / factor) * \
           (1 - np.abs(og[1] - center) / factor)
    # Flax uses HWIO format for ConvTranspose kernels
    weight = np.zeros((kernel_size, kernel_size, in_channels, out_channels))
    for i in range(min(in_channels, out_channels)):
        weight[:, :, i, i] = filt
    return jnp.array(weight)
def bilinear_kernel(in_channels, out_channels, kernel_size):
    factor = (kernel_size + 1) // 2
    if kernel_size % 2 == 1:
        center = factor - 1
    else:
        center = factor - 0.5
    og = (np.arange(kernel_size).reshape(-1, 1),
          np.arange(kernel_size).reshape(1, -1))
    filt = (1 - np.abs(og[0] - center) / factor) * \
           (1 - np.abs(og[1] - center) / factor)
    weight = np.zeros((in_channels, out_channels, kernel_size, kernel_size))
    weight[range(in_channels), range(out_channels), :, :] = filt
    return np.array(weight)

Let’s experiment with upsampling of bilinear interpolation that is implemented by a transposed convolutional layer. We construct a transposed convolutional layer that doubles the height and weight, and initialize its kernel with the bilinear_kernel function.

conv_trans = nn.ConvTranspose2d(3, 3, kernel_size=4, padding=1, stride=2,
                                bias=False)
with torch.no_grad():
    conv_trans.weight.copy_(bilinear_kernel(3, 3, 4));
# Build a transposed conv layer with bilinear initialization to double H and W
bilinear_w = bilinear_kernel(3, 3, 4)
conv_trans = keras.layers.Conv2DTranspose(
    3, kernel_size=4, strides=2, padding='same', use_bias=False,
    kernel_initializer=tf.constant_initializer(bilinear_w))
# Build the layer by passing a dummy input
_ = conv_trans(tf.zeros((1, 1, 1, 3)))
class BilinearConvTranspose(nnx.Module):
    """A transposed conv layer initialized with bilinear interpolation."""
    def __init__(self, channels, kernel_size, strides, *, rngs):
        self.layer = nnx.ConvTranspose(
            channels, channels, (kernel_size, kernel_size), strides=strides,
            padding='SAME', use_bias=False, rngs=rngs)
        self.layer.kernel[...] = bilinear_kernel(
            channels, channels, kernel_size)

    def __call__(self, X):
        return self.layer(X)

conv_trans = BilinearConvTranspose(3, 4, (2, 2), rngs=nnx.Rngs(0))
conv_trans = nn.Conv2DTranspose(3, kernel_size=4, padding=1, strides=2)
conv_trans.initialize(init.Constant(bilinear_kernel(3, 3, 4)))

Read the image X and assign the upsampling output to Y. In order to print the image, we need to adjust the position of the channel dimension.

img = torchvision.transforms.ToTensor()(d2l.Image.open('../img/catdog.jpg'))
X = img.unsqueeze(0)
Y = conv_trans(X)
out_img = Y[0].permute(1, 2, 0).detach()
img = np.array(Image.open('../img/catdog.jpg')).astype(np.float32) / 255
X = tf.expand_dims(tf.constant(img), axis=0)  # NHWC
Y = conv_trans(X)
out_img = Y[0].numpy()
img = np.array(Image.open('../img/catdog.jpg')).astype(np.float32) / 255
X = jnp.expand_dims(jnp.array(img), axis=0)  # NHWC
Y = conv_trans(X)
out_img = np.array(Y[0])
E0719 14:24:24.060179   40677 cuda_timer.cc:88] Delay kernel timed out: measured time has sub-optimal accuracy. There may be a missing warmup execution, please investigate in Nsight Systems.
img = image.imread('../img/catdog.jpg')
X = np.expand_dims(img.astype('float32').transpose(2, 0, 1), axis=0) / 255
Y = conv_trans(X)
out_img = Y[0].transpose(1, 2, 0)

As we can see, the transposed convolutional layer increases both the height and width of the image by a factor of two. Except for the different scales in coordinates, the image scaled up by bilinear interpolation and the original image printed in Section 19.3 look the same.

d2l.set_figsize()
print('input image shape:', img.permute(1, 2, 0).shape)
d2l.plt.imshow(img.permute(1, 2, 0));
print('output image shape:', out_img.shape)
d2l.plt.imshow(out_img);
input image shape: torch.Size([561, 728, 3])
output image shape: torch.Size([1122, 1456, 3])

d2l.set_figsize()
print('input image shape:', img.shape)
d2l.plt.imshow(img);
print('output image shape:', out_img.shape)
d2l.plt.imshow(np.clip(out_img, 0, 1));
input image shape: (561, 728, 3)
output image shape: (1122, 1456, 3)

d2l.set_figsize()
print('input image shape:', img.shape)
d2l.plt.imshow(img);
print('output image shape:', out_img.shape)
d2l.plt.imshow(out_img);
input image shape: (561, 728, 3)
output image shape: (1122, 1456, 3)

d2l.set_figsize()
print('input image shape:', img.shape)
d2l.plt.imshow(img.asnumpy());
print('output image shape:', out_img.shape)
d2l.plt.imshow(out_img.asnumpy());
input image shape: (561, 728, 3)
output image shape: (1122, 1456, 3)

In a fully convolutional network, we initialize the transposed convolutional layer with upsampling of bilinear interpolation. For the \(1\times 1\) convolutional layer, we use Xavier initialization.

W = bilinear_kernel(num_classes, num_classes, 64)
net.transpose_conv.weight.data.copy_(W);
# Initialize the transpose conv kernel with bilinear upsampling weights.
# The 1x1 conv was already initialized with Glorot (Xavier) uniform above.
W = bilinear_kernel(num_classes, num_classes, 64)
# Find the Conv2DTranspose layer in fcn_net and set its weights
for layer in fcn_net.layers:
    if isinstance(layer, keras.layers.Conv2DTranspose):
        layer.set_weights([W])
        break
# Initialize the FCN with bilinear weights for the transposed conv layer
# and Xavier initialization for the 1x1 conv layer
W = bilinear_kernel(num_classes, num_classes, 64)
net.upsample.kernel[...] = W
W = bilinear_kernel(num_classes, num_classes, 64)
net[-1].initialize(init.Constant(W))
net[-2].initialize(init=init.Xavier())

19.11.3 Reading the Dataset

We read the semantic segmentation dataset as introduced in Section 19.9. The output image shape of random cropping is specified as \(320\times 480\): both the height and width are divisible by \(32\).

batch_size, crop_size = 32, (320, 480)
train_iter, test_iter = d2l.load_data_voc(batch_size, crop_size)
read 1114 examples
read 1078 examples
batch_size, crop_size = 32, (320, 480)
train_iter, test_iter = d2l.load_data_voc(batch_size, crop_size)
read 1114 examples
read 1078 examples
# The NNX ResNet-50 is deeper than the ResNet-18 used in the PyTorch tab, so
# use a smaller minibatch while keeping the same crops and number of epochs.
batch_size, crop_size = 4, (320, 480)
train_iter, test_iter = d2l.load_data_voc(batch_size, crop_size)
read 1114 examples
read 1078 examples
batch_size, crop_size = 32, (320, 480)
train_iter, test_iter = d2l.load_data_voc(batch_size, crop_size)
read 1114 examples
read 1078 examples

19.11.4 Training

Now we can train our constructed fully convolutional network. The loss function and accuracy calculation here are not essentially different from those in image classification of earlier chapters. Because we use the output channel of the transposed convolutional layer to predict the class for each pixel, the channel dimension is specified in the loss calculation. In addition, the accuracy is calculated based on correctness of the predicted class for all the pixels.

def loss(inputs, targets):
    return F.cross_entropy(inputs, targets, reduction='none').mean(1).mean(1)

num_epochs, lr, wd, devices = 5, 0.001, 1e-3, d2l.try_all_gpus()
trainer = torch.optim.SGD(net.parameters(), lr=lr, weight_decay=wd)
d2l.train_ch13(net, train_iter, test_iter, loss, trainer, num_epochs, devices)
loss 0.411, train acc 0.870, test acc 0.851
351.0 examples/sec on [device(type='cuda', index=0)]

# Loss: SparseCategoricalCrossentropy over per-pixel logits (NHWC -> NHW).
# Full fine-tuning of the entire network (backbone + head) to match the
# PyTorch tab.
num_epochs, lr, wd = 5, 0.001, 1e-3
fcn_net.compile(
    optimizer=keras.optimizers.SGD(learning_rate=lr, weight_decay=wd),
    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=['accuracy'],
    # Keras may otherwise select XLA automatically. For this unusually large
    # transposed convolution, XLA's cuDNN autotuning can take tens of minutes
    # before the first training step and offers no benefit to this example.
    jit_compile=False)
fcn_net.fit(train_iter, epochs=num_epochs, validation_data=test_iter)
Epoch 1/5
 1/34 ━━━━━━━━━━━━━━━━━━━━ 18:08 33s/step - accuracy: 0.0158 - loss: 4.1711
 2/34 ━━━━━━━━━━━━━━━━━━━━ 5s 176ms/step - accuracy: 0.0141 - loss: 4.1660 
 3/34 ━━━━━━━━━━━━━━━━━━━━ 5s 169ms/step - accuracy: 0.0136 - loss: 4.1353
 4/34 ━━━━━━━━━━━━━━━━━━━━ 5s 172ms/step - accuracy: 0.0135 - loss: 4.1040
 5/34 ━━━━━━━━━━━━━━━━━━━━ 4s 164ms/step - accuracy: 0.0140 - loss: 4.0669
 6/34 ━━━━━━━━━━━━━━━━━━━━ 4s 166ms/step - accuracy: 0.0148 - loss: 4.0280
 7/34 ━━━━━━━━━━━━━━━━━━━━ 4s 171ms/step - accuracy: 0.0162 - loss: 3.9916
 8/34 ━━━━━━━━━━━━━━━━━━━━ 4s 172ms/step - accuracy: 0.0184 - loss: 3.9551
 9/34 ━━━━━━━━━━━━━━━━━━━━ 4s 169ms/step - accuracy: 0.0212 - loss: 3.9183
10/34 ━━━━━━━━━━━━━━━━━━━━ 3s 162ms/step - accuracy: 0.0250 - loss: 3.8806
11/34 ━━━━━━━━━━━━━━━━━━━━ 3s 161ms/step - accuracy: 0.0297 - loss: 3.8422
12/34 ━━━━━━━━━━━━━━━━━━━━ 3s 161ms/step - accuracy: 0.0350 - loss: 3.8044
13/34 ━━━━━━━━━━━━━━━━━━━━ 3s 162ms/step - accuracy: 0.0412 - loss: 3.7666
14/34 ━━━━━━━━━━━━━━━━━━━━ 3s 161ms/step - accuracy: 0.0480 - loss: 3.7296
15/34 ━━━━━━━━━━━━━━━━━━━━ 3s 161ms/step - accuracy: 0.0552 - loss: 3.6937
16/34 ━━━━━━━━━━━━━━━━━━━━ 2s 159ms/step - accuracy: 0.0630 - loss: 3.6578
17/34 ━━━━━━━━━━━━━━━━━━━━ 2s 158ms/step - accuracy: 0.0712 - loss: 3.6226
18/34 ━━━━━━━━━━━━━━━━━━━━ 2s 160ms/step - accuracy: 0.0794 - loss: 3.5891
19/34 ━━━━━━━━━━━━━━━━━━━━ 2s 162ms/step - accuracy: 0.0878 - loss: 3.5562
20/34 ━━━━━━━━━━━━━━━━━━━━ 2s 183ms/step - accuracy: 0.0963 - loss: 3.5246
21/34 ━━━━━━━━━━━━━━━━━━━━ 2s 181ms/step - accuracy: 0.1047 - loss: 3.4940
22/34 ━━━━━━━━━━━━━━━━━━━━ 2s 181ms/step - accuracy: 0.1131 - loss: 3.4640
23/34 ━━━━━━━━━━━━━━━━━━━━ 1s 180ms/step - accuracy: 0.1215 - loss: 3.4346
24/34 ━━━━━━━━━━━━━━━━━━━━ 1s 180ms/step - accuracy: 0.1298 - loss: 3.4061
25/34 ━━━━━━━━━━━━━━━━━━━━ 1s 178ms/step - accuracy: 0.1379 - loss: 3.3790
26/34 ━━━━━━━━━━━━━━━━━━━━ 1s 177ms/step - accuracy: 0.1458 - loss: 3.3525
27/34 ━━━━━━━━━━━━━━━━━━━━ 1s 175ms/step - accuracy: 0.1537 - loss: 3.3267
28/34 ━━━━━━━━━━━━━━━━━━━━ 1s 174ms/step - accuracy: 0.1614 - loss: 3.3016
29/34 ━━━━━━━━━━━━━━━━━━━━ 0s 173ms/step - accuracy: 0.1689 - loss: 3.2776
30/34 ━━━━━━━━━━━━━━━━━━━━ 0s 170ms/step - accuracy: 0.1763 - loss: 3.2540
31/34 ━━━━━━━━━━━━━━━━━━━━ 0s 168ms/step - accuracy: 0.1834 - loss: 3.2313
32/34 ━━━━━━━━━━━━━━━━━━━━ 0s 166ms/step - accuracy: 0.1904 - loss: 3.2092
33/34 ━━━━━━━━━━━━━━━━━━━━ 0s 164ms/step - accuracy: 0.1973 - loss: 3.1877
34/34 ━━━━━━━━━━━━━━━━━━━━ 0s 162ms/step - accuracy: 0.2039 - loss: 3.1668
34/34 ━━━━━━━━━━━━━━━━━━━━ 46s 390ms/step - accuracy: 0.4238 - loss: 2.4785 - val_accuracy: 0.7288 - val_loss: 2.5190
Epoch 2/5
 1/34 ━━━━━━━━━━━━━━━━━━━━ 9s 291ms/step - accuracy: 0.6792 - loss: 1.7432
 2/34 ━━━━━━━━━━━━━━━━━━━━ 12s 390ms/step - accuracy: 0.6964 - loss: 1.6723
 3/34 ━━━━━━━━━━━━━━━━━━━━ 9s 317ms/step - accuracy: 0.7042 - loss: 1.6489 
 4/34 ━━━━━━━━━━━━━━━━━━━━ 8s 293ms/step - accuracy: 0.7065 - loss: 1.6438
 5/34 ━━━━━━━━━━━━━━━━━━━━ 7s 270ms/step - accuracy: 0.7083 - loss: 1.6380
 6/34 ━━━━━━━━━━━━━━━━━━━━ 7s 261ms/step - accuracy: 0.7076 - loss: 1.6405
 7/34 ━━━━━━━━━━━━━━━━━━━━ 6s 243ms/step - accuracy: 0.7076 - loss: 1.6398
 8/34 ━━━━━━━━━━━━━━━━━━━━ 6s 239ms/step - accuracy: 0.7078 - loss: 1.6396
 9/34 ━━━━━━━━━━━━━━━━━━━━ 5s 231ms/step - accuracy: 0.7085 - loss: 1.6367
10/34 ━━━━━━━━━━━━━━━━━━━━ 5s 223ms/step - accuracy: 0.7088 - loss: 1.6349
11/34 ━━━━━━━━━━━━━━━━━━━━ 4s 216ms/step - accuracy: 0.7089 - loss: 1.6344
12/34 ━━━━━━━━━━━━━━━━━━━━ 4s 206ms/step - accuracy: 0.7093 - loss: 1.6326
13/34 ━━━━━━━━━━━━━━━━━━━━ 4s 202ms/step - accuracy: 0.7097 - loss: 1.6308
14/34 ━━━━━━━━━━━━━━━━━━━━ 3s 199ms/step - accuracy: 0.7103 - loss: 1.6281
15/34 ━━━━━━━━━━━━━━━━━━━━ 3s 193ms/step - accuracy: 0.7108 - loss: 1.6259
16/34 ━━━━━━━━━━━━━━━━━━━━ 3s 190ms/step - accuracy: 0.7111 - loss: 1.6242
17/34 ━━━━━━━━━━━━━━━━━━━━ 3s 187ms/step - accuracy: 0.7114 - loss: 1.6222
18/34 ━━━━━━━━━━━━━━━━━━━━ 2s 186ms/step - accuracy: 0.7114 - loss: 1.6222
19/34 ━━━━━━━━━━━━━━━━━━━━ 2s 185ms/step - accuracy: 0.7112 - loss: 1.6227
20/34 ━━━━━━━━━━━━━━━━━━━━ 2s 184ms/step - accuracy: 0.7110 - loss: 1.6233
21/34 ━━━━━━━━━━━━━━━━━━━━ 2s 182ms/step - accuracy: 0.7108 - loss: 1.6237
22/34 ━━━━━━━━━━━━━━━━━━━━ 2s 182ms/step - accuracy: 0.7107 - loss: 1.6241
23/34 ━━━━━━━━━━━━━━━━━━━━ 1s 180ms/step - accuracy: 0.7105 - loss: 1.6243
24/34 ━━━━━━━━━━━━━━━━━━━━ 1s 180ms/step - accuracy: 0.7103 - loss: 1.6243
25/34 ━━━━━━━━━━━━━━━━━━━━ 1s 179ms/step - accuracy: 0.7102 - loss: 1.6246
26/34 ━━━━━━━━━━━━━━━━━━━━ 1s 179ms/step - accuracy: 0.7101 - loss: 1.6242
27/34 ━━━━━━━━━━━━━━━━━━━━ 1s 178ms/step - accuracy: 0.7101 - loss: 1.6239
28/34 ━━━━━━━━━━━━━━━━━━━━ 1s 175ms/step - accuracy: 0.7101 - loss: 1.6233
29/34 ━━━━━━━━━━━━━━━━━━━━ 0s 173ms/step - accuracy: 0.7100 - loss: 1.6228
30/34 ━━━━━━━━━━━━━━━━━━━━ 0s 170ms/step - accuracy: 0.7100 - loss: 1.6223
31/34 ━━━━━━━━━━━━━━━━━━━━ 0s 168ms/step - accuracy: 0.7099 - loss: 1.6219
32/34 ━━━━━━━━━━━━━━━━━━━━ 0s 166ms/step - accuracy: 0.7099 - loss: 1.6214
33/34 ━━━━━━━━━━━━━━━━━━━━ 0s 163ms/step - accuracy: 0.7098 - loss: 1.6211
34/34 ━━━━━━━━━━━━━━━━━━━━ 0s 161ms/step - accuracy: 0.7097 - loss: 1.6207
34/34 ━━━━━━━━━━━━━━━━━━━━ 10s 286ms/step - accuracy: 0.7071 - loss: 1.6097 - val_accuracy: 0.7285 - val_loss: 1.9542
Epoch 3/5
 1/34 ━━━━━━━━━━━━━━━━━━━━ 8s 256ms/step - accuracy: 0.7457 - loss: 1.3850
 2/34 ━━━━━━━━━━━━━━━━━━━━ 10s 313ms/step - accuracy: 0.7336 - loss: 1.4420
 3/34 ━━━━━━━━━━━━━━━━━━━━ 9s 293ms/step - accuracy: 0.7231 - loss: 1.4825 
 4/34 ━━━━━━━━━━━━━━━━━━━━ 7s 266ms/step - accuracy: 0.7187 - loss: 1.5014
 5/34 ━━━━━━━━━━━━━━━━━━━━ 8s 288ms/step - accuracy: 0.7193 - loss: 1.4988
 6/34 ━━━━━━━━━━━━━━━━━━━━ 7s 262ms/step - accuracy: 0.7203 - loss: 1.4952
 7/34 ━━━━━━━━━━━━━━━━━━━━ 6s 244ms/step - accuracy: 0.7206 - loss: 1.4958
 8/34 ━━━━━━━━━━━━━━━━━━━━ 6s 231ms/step - accuracy: 0.7200 - loss: 1.4988
 9/34 ━━━━━━━━━━━━━━━━━━━━ 5s 222ms/step - accuracy: 0.7193 - loss: 1.5026
10/34 ━━━━━━━━━━━━━━━━━━━━ 5s 213ms/step - accuracy: 0.7190 - loss: 1.5040
11/34 ━━━━━━━━━━━━━━━━━━━━ 4s 207ms/step - accuracy: 0.7188 - loss: 1.5050
12/34 ━━━━━━━━━━━━━━━━━━━━ 4s 200ms/step - accuracy: 0.7186 - loss: 1.5059
13/34 ━━━━━━━━━━━━━━━━━━━━ 4s 196ms/step - accuracy: 0.7188 - loss: 1.5051
14/34 ━━━━━━━━━━━━━━━━━━━━ 3s 191ms/step - accuracy: 0.7187 - loss: 1.5045
15/34 ━━━━━━━━━━━━━━━━━━━━ 3s 188ms/step - accuracy: 0.7186 - loss: 1.5047
16/34 ━━━━━━━━━━━━━━━━━━━━ 3s 185ms/step - accuracy: 0.7185 - loss: 1.5045
17/34 ━━━━━━━━━━━━━━━━━━━━ 3s 184ms/step - accuracy: 0.7183 - loss: 1.5055
18/34 ━━━━━━━━━━━━━━━━━━━━ 2s 182ms/step - accuracy: 0.7178 - loss: 1.5071
19/34 ━━━━━━━━━━━━━━━━━━━━ 2s 179ms/step - accuracy: 0.7174 - loss: 1.5082
20/34 ━━━━━━━━━━━━━━━━━━━━ 2s 179ms/step - accuracy: 0.7171 - loss: 1.5093
21/34 ━━━━━━━━━━━━━━━━━━━━ 2s 178ms/step - accuracy: 0.7168 - loss: 1.5102
22/34 ━━━━━━━━━━━━━━━━━━━━ 2s 176ms/step - accuracy: 0.7165 - loss: 1.5108
23/34 ━━━━━━━━━━━━━━━━━━━━ 1s 175ms/step - accuracy: 0.7162 - loss: 1.5116
24/34 ━━━━━━━━━━━━━━━━━━━━ 1s 175ms/step - accuracy: 0.7159 - loss: 1.5123
25/34 ━━━━━━━━━━━━━━━━━━━━ 1s 175ms/step - accuracy: 0.7156 - loss: 1.5133
26/34 ━━━━━━━━━━━━━━━━━━━━ 1s 175ms/step - accuracy: 0.7153 - loss: 1.5139
27/34 ━━━━━━━━━━━━━━━━━━━━ 1s 173ms/step - accuracy: 0.7152 - loss: 1.5138
28/34 ━━━━━━━━━━━━━━━━━━━━ 1s 172ms/step - accuracy: 0.7151 - loss: 1.5141
29/34 ━━━━━━━━━━━━━━━━━━━━ 0s 169ms/step - accuracy: 0.7149 - loss: 1.5145
30/34 ━━━━━━━━━━━━━━━━━━━━ 0s 167ms/step - accuracy: 0.7148 - loss: 1.5149
31/34 ━━━━━━━━━━━━━━━━━━━━ 0s 164ms/step - accuracy: 0.7146 - loss: 1.5151
32/34 ━━━━━━━━━━━━━━━━━━━━ 0s 162ms/step - accuracy: 0.7145 - loss: 1.5155
33/34 ━━━━━━━━━━━━━━━━━━━━ 0s 160ms/step - accuracy: 0.7143 - loss: 1.5158
34/34 ━━━━━━━━━━━━━━━━━━━━ 0s 158ms/step - accuracy: 0.7142 - loss: 1.5161
34/34 ━━━━━━━━━━━━━━━━━━━━ 10s 280ms/step - accuracy: 0.7097 - loss: 1.5249 - val_accuracy: 0.7294 - val_loss: 1.6552
Epoch 4/5
 1/34 ━━━━━━━━━━━━━━━━━━━━ 8s 249ms/step - accuracy: 0.7969 - loss: 1.1104
 2/34 ━━━━━━━━━━━━━━━━━━━━ 14s 461ms/step - accuracy: 0.7941 - loss: 1.1287
 3/34 ━━━━━━━━━━━━━━━━━━━━ 11s 373ms/step - accuracy: 0.7838 - loss: 1.1779
 4/34 ━━━━━━━━━━━━━━━━━━━━ 9s 321ms/step - accuracy: 0.7751 - loss: 1.2138 
 5/34 ━━━━━━━━━━━━━━━━━━━━ 8s 278ms/step - accuracy: 0.7672 - loss: 1.2459
 6/34 ━━━━━━━━━━━━━━━━━━━━ 6s 245ms/step - accuracy: 0.7616 - loss: 1.2682
 7/34 ━━━━━━━━━━━━━━━━━━━━ 6s 234ms/step - accuracy: 0.7582 - loss: 1.2813
 8/34 ━━━━━━━━━━━━━━━━━━━━ 5s 220ms/step - accuracy: 0.7549 - loss: 1.2935
 9/34 ━━━━━━━━━━━━━━━━━━━━ 5s 213ms/step - accuracy: 0.7524 - loss: 1.3023
10/34 ━━━━━━━━━━━━━━━━━━━━ 4s 207ms/step - accuracy: 0.7504 - loss: 1.3098
11/34 ━━━━━━━━━━━━━━━━━━━━ 4s 205ms/step - accuracy: 0.7485 - loss: 1.3175
12/34 ━━━━━━━━━━━━━━━━━━━━ 4s 201ms/step - accuracy: 0.7465 - loss: 1.3259
13/34 ━━━━━━━━━━━━━━━━━━━━ 4s 196ms/step - accuracy: 0.7447 - loss: 1.3337
14/34 ━━━━━━━━━━━━━━━━━━━━ 3s 194ms/step - accuracy: 0.7428 - loss: 1.3418
15/34 ━━━━━━━━━━━━━━━━━━━━ 3s 190ms/step - accuracy: 0.7412 - loss: 1.3489
16/34 ━━━━━━━━━━━━━━━━━━━━ 3s 189ms/step - accuracy: 0.7398 - loss: 1.3550
17/34 ━━━━━━━━━━━━━━━━━━━━ 3s 188ms/step - accuracy: 0.7385 - loss: 1.3606
18/34 ━━━━━━━━━━━━━━━━━━━━ 2s 187ms/step - accuracy: 0.7372 - loss: 1.3660
19/34 ━━━━━━━━━━━━━━━━━━━━ 2s 186ms/step - accuracy: 0.7360 - loss: 1.3709
20/34 ━━━━━━━━━━━━━━━━━━━━ 2s 184ms/step - accuracy: 0.7350 - loss: 1.3747
21/34 ━━━━━━━━━━━━━━━━━━━━ 2s 183ms/step - accuracy: 0.7341 - loss: 1.3785
22/34 ━━━━━━━━━━━━━━━━━━━━ 2s 182ms/step - accuracy: 0.7334 - loss: 1.3815
23/34 ━━━━━━━━━━━━━━━━━━━━ 1s 182ms/step - accuracy: 0.7326 - loss: 1.3848
24/34 ━━━━━━━━━━━━━━━━━━━━ 1s 179ms/step - accuracy: 0.7319 - loss: 1.3880
25/34 ━━━━━━━━━━━━━━━━━━━━ 1s 179ms/step - accuracy: 0.7312 - loss: 1.3910
26/34 ━━━━━━━━━━━━━━━━━━━━ 1s 178ms/step - accuracy: 0.7305 - loss: 1.3939
27/34 ━━━━━━━━━━━━━━━━━━━━ 1s 176ms/step - accuracy: 0.7298 - loss: 1.3965
28/34 ━━━━━━━━━━━━━━━━━━━━ 1s 174ms/step - accuracy: 0.7292 - loss: 1.3990
29/34 ━━━━━━━━━━━━━━━━━━━━ 0s 171ms/step - accuracy: 0.7287 - loss: 1.4013
30/34 ━━━━━━━━━━━━━━━━━━━━ 0s 169ms/step - accuracy: 0.7282 - loss: 1.4032
31/34 ━━━━━━━━━━━━━━━━━━━━ 0s 167ms/step - accuracy: 0.7277 - loss: 1.4049
32/34 ━━━━━━━━━━━━━━━━━━━━ 0s 165ms/step - accuracy: 0.7273 - loss: 1.4065
33/34 ━━━━━━━━━━━━━━━━━━━━ 0s 162ms/step - accuracy: 0.7270 - loss: 1.4080
34/34 ━━━━━━━━━━━━━━━━━━━━ 0s 160ms/step - accuracy: 0.7266 - loss: 1.4095
34/34 ━━━━━━━━━━━━━━━━━━━━ 10s 280ms/step - accuracy: 0.7142 - loss: 1.4586 - val_accuracy: 0.7279 - val_loss: 1.5219
Epoch 5/5
 1/34 ━━━━━━━━━━━━━━━━━━━━ 9s 278ms/step - accuracy: 0.7326 - loss: 1.3792
 2/34 ━━━━━━━━━━━━━━━━━━━━ 9s 312ms/step - accuracy: 0.7365 - loss: 1.3567
 3/34 ━━━━━━━━━━━━━━━━━━━━ 8s 282ms/step - accuracy: 0.7265 - loss: 1.3870
 4/34 ━━━━━━━━━━━━━━━━━━━━ 9s 303ms/step - accuracy: 0.7205 - loss: 1.4067
 5/34 ━━━━━━━━━━━━━━━━━━━━ 7s 267ms/step - accuracy: 0.7168 - loss: 1.4206
 6/34 ━━━━━━━━━━━━━━━━━━━━ 6s 241ms/step - accuracy: 0.7159 - loss: 1.4226
 7/34 ━━━━━━━━━━━━━━━━━━━━ 6s 228ms/step - accuracy: 0.7158 - loss: 1.4219
 8/34 ━━━━━━━━━━━━━━━━━━━━ 5s 219ms/step - accuracy: 0.7169 - loss: 1.4164
 9/34 ━━━━━━━━━━━━━━━━━━━━ 5s 206ms/step - accuracy: 0.7171 - loss: 1.4155
10/34 ━━━━━━━━━━━━━━━━━━━━ 4s 206ms/step - accuracy: 0.7168 - loss: 1.4170
11/34 ━━━━━━━━━━━━━━━━━━━━ 4s 202ms/step - accuracy: 0.7169 - loss: 1.4163
12/34 ━━━━━━━━━━━━━━━━━━━━ 4s 197ms/step - accuracy: 0.7172 - loss: 1.4150
13/34 ━━━━━━━━━━━━━━━━━━━━ 4s 195ms/step - accuracy: 0.7176 - loss: 1.4130
14/34 ━━━━━━━━━━━━━━━━━━━━ 3s 190ms/step - accuracy: 0.7181 - loss: 1.4108
15/34 ━━━━━━━━━━━━━━━━━━━━ 3s 187ms/step - accuracy: 0.7186 - loss: 1.4083
16/34 ━━━━━━━━━━━━━━━━━━━━ 3s 187ms/step - accuracy: 0.7188 - loss: 1.4075
17/34 ━━━━━━━━━━━━━━━━━━━━ 3s 184ms/step - accuracy: 0.7187 - loss: 1.4077
18/34 ━━━━━━━━━━━━━━━━━━━━ 2s 183ms/step - accuracy: 0.7188 - loss: 1.4075
19/34 ━━━━━━━━━━━━━━━━━━━━ 2s 181ms/step - accuracy: 0.7190 - loss: 1.4068
20/34 ━━━━━━━━━━━━━━━━━━━━ 2s 177ms/step - accuracy: 0.7191 - loss: 1.4061
21/34 ━━━━━━━━━━━━━━━━━━━━ 2s 176ms/step - accuracy: 0.7193 - loss: 1.4049
22/34 ━━━━━━━━━━━━━━━━━━━━ 2s 175ms/step - accuracy: 0.7193 - loss: 1.4045
23/34 ━━━━━━━━━━━━━━━━━━━━ 1s 176ms/step - accuracy: 0.7193 - loss: 1.4046
24/34 ━━━━━━━━━━━━━━━━━━━━ 1s 173ms/step - accuracy: 0.7191 - loss: 1.4051
25/34 ━━━━━━━━━━━━━━━━━━━━ 1s 173ms/step - accuracy: 0.7190 - loss: 1.4055
26/34 ━━━━━━━━━━━━━━━━━━━━ 1s 173ms/step - accuracy: 0.7189 - loss: 1.4057
27/34 ━━━━━━━━━━━━━━━━━━━━ 1s 173ms/step - accuracy: 0.7187 - loss: 1.4061
28/34 ━━━━━━━━━━━━━━━━━━━━ 1s 172ms/step - accuracy: 0.7186 - loss: 1.4067
29/34 ━━━━━━━━━━━━━━━━━━━━ 0s 169ms/step - accuracy: 0.7184 - loss: 1.4073
30/34 ━━━━━━━━━━━━━━━━━━━━ 0s 167ms/step - accuracy: 0.7182 - loss: 1.4080
31/34 ━━━━━━━━━━━━━━━━━━━━ 0s 165ms/step - accuracy: 0.7180 - loss: 1.4085
32/34 ━━━━━━━━━━━━━━━━━━━━ 0s 163ms/step - accuracy: 0.7179 - loss: 1.4088
33/34 ━━━━━━━━━━━━━━━━━━━━ 0s 161ms/step - accuracy: 0.7178 - loss: 1.4093
34/34 ━━━━━━━━━━━━━━━━━━━━ 0s 159ms/step - accuracy: 0.7176 - loss: 1.4100
34/34 ━━━━━━━━━━━━━━━━━━━━ 9s 277ms/step - accuracy: 0.7120 - loss: 1.4347 - val_accuracy: 0.7288 - val_loss: 1.4652
<keras.src.callbacks.history.History at 0x7722ac22b830>
num_epochs, lr, wd = 5, 0.001, 1e-3
def parameter_labels(params):
    return jax.tree_util.tree_map_with_path(
        lambda path, _: ('head' if any(
            getattr(p, 'key', None) in ('classifier', 'upsample')
            for p in path) else 'backbone'), params)

optimizer = nnx.Optimizer(
    net, optax.multi_transform(
        {'head': optax.chain(optax.add_decayed_weights(wd),
                             optax.sgd(lr * 10)),
         'backbone': optax.chain(optax.add_decayed_weights(wd),
                                 optax.sgd(lr))},
        parameter_labels),
    wrt=nnx.Param)
eval_net = nnx.view(net, use_running_average=True, raise_if_not_found=False)

@nnx.jit
def train_step(model, optimizer, X, Y):
    def loss_fn(model):
        logits = model(X)
        loss = optax.softmax_cross_entropy_with_integer_labels(logits, Y)
        return loss.mean(), logits
    (loss, logits), grads = nnx.value_and_grad(loss_fn, has_aux=True)(model)
    optimizer.update(model, grads)
    correct = (jnp.argmax(logits, axis=-1) == Y).sum()
    return loss, correct

@nnx.jit
def eval_step(model, X, Y):
    logits = model(X)
    losses = optax.softmax_cross_entropy_with_integer_labels(logits, Y)
    correct = (jnp.argmax(logits, axis=-1) == Y).sum()
    return losses.sum(), correct

for epoch in range(num_epochs):
    loss_terms, correct_terms, num_pixels = [], [], 0
    for X, Y in train_iter:
        X = jnp.transpose(jnp.array(X), (0, 2, 3, 1))
        Y = jnp.array(Y)
        loss, correct = train_step(net, optimizer, X, Y)
        loss_terms.append(loss * Y.size)
        correct_terms.append(correct)
        num_pixels += Y.size
    val_loss_terms, val_correct_terms, val_pixels = [], [], 0
    for X, Y in test_iter:
        X = jnp.transpose(jnp.array(X), (0, 2, 3, 1))
        Y = jnp.array(Y)
        val_loss, val_correct = eval_step(eval_net, X, Y)
        val_loss_terms.append(val_loss)
        val_correct_terms.append(val_correct)
        val_pixels += Y.size
    loss_sum = float(jnp.stack(loss_terms).sum())
    correct = int(jnp.stack(correct_terms).sum())
    val_loss_sum = float(jnp.stack(val_loss_terms).sum())
    val_correct = int(jnp.stack(val_correct_terms).sum())
    print(f'epoch {epoch + 1}, loss {loss_sum / num_pixels:.3f}, '
          f'pixel acc {correct / num_pixels:.3f}, '
          f'val loss {val_loss_sum / val_pixels:.3f}, '
          f'val pixel acc {val_correct / val_pixels:.3f}')
epoch 1, loss 1.635, pixel acc 0.706, val loss 1.202, val pixel acc 0.731
epoch 2, loss 1.132, pixel acc 0.738, val loss 1.059, val pixel acc 0.747
epoch 3, loss 1.016, pixel acc 0.757, val loss 0.974, val pixel acc 0.765
epoch 4, loss 0.924, pixel acc 0.775, val loss 0.886, val pixel acc 0.781
epoch 5, loss 0.852, pixel acc 0.790, val loss 0.826, val pixel acc 0.790
num_epochs, lr, wd, devices = 5, 0.001, 1e-3, d2l.try_all_gpus()
loss = gluon.loss.SoftmaxCrossEntropyLoss(axis=1)
net.reset_ctx(devices)
trainer = gluon.Trainer(net.collect_params(), 'sgd',
                        {'learning_rate': lr, 'wd': wd})
d2l.train_ch13(net, train_iter, test_iter, loss, trainer, num_epochs, devices)
loss 0.370, train acc 0.881, test acc 0.846
85.1 examples/sec on [gpu(0)]

19.11.5 Prediction

When predicting, we need to standardize the input image in each channel and transform the image into the four-dimensional input format required by the CNN.

def predict(img):
    net.eval()
    X = test_iter.dataset.normalize_image(img).unsqueeze(0)
    with torch.no_grad():
        pred = net(X.to(devices[0])).argmax(dim=1)
    return pred.reshape(pred.shape[1], pred.shape[2])
def predict(img):
    rgb_mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
    rgb_std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
    X = (img.astype(np.float32) / 255 - rgb_mean) / rgb_std
    X = tf.expand_dims(tf.constant(X), axis=0)  # NHWC
    pred = fcn_net(X, training=False)  # (1, H, W, num_classes)
    return tf.reshape(tf.argmax(pred, axis=-1), pred.shape[1:3])
def predict(img):
    rgb_mean = np.array([0.485, 0.456, 0.406])
    rgb_std = np.array([0.229, 0.224, 0.225])
    X = (img.astype(np.float32) / 255 - rgb_mean) / rgb_std
    X = jnp.expand_dims(jnp.array(X), axis=0)  # NHWC
    pred = nnx.view(net, use_running_average=True,
                    raise_if_not_found=False)(X)
    return jnp.argmax(pred, axis=-1).reshape(pred.shape[1], pred.shape[2])
def predict(img):
    X = test_iter._dataset.normalize_image(img)
    X = np.expand_dims(X.transpose(2, 0, 1), axis=0)
    pred = net(X.as_in_ctx(devices[0])).argmax(axis=1)
    return pred.reshape(pred.shape[1], pred.shape[2])

To visualize the predicted class of each pixel, we map the predicted class back to its label color in the dataset.

def label2image(pred):
    colormap = torch.tensor(d2l.VOC_COLORMAP, device=devices[0])
    X = pred.long()
    return colormap[X, :]
def label2image(pred):
    colormap = tf.constant(d2l.VOC_COLORMAP, dtype=tf.uint8)
    X = tf.cast(pred, tf.int32)
    return tf.gather(colormap, X)
def label2image(pred):
    colormap = jnp.array(d2l.VOC_COLORMAP, dtype=jnp.uint8)
    X = pred.astype(jnp.int32)
    return colormap[X, :]
def label2image(pred):
    colormap = np.array(d2l.VOC_COLORMAP, ctx=devices[0], dtype='uint8')
    X = pred.astype('int32')
    return colormap[X, :]

Images in the test dataset vary in size and shape. Since the model uses a transposed convolutional layer with stride of 32, when the height or width of an input image is indivisible by 32, the output height or width of the transposed convolutional layer will deviate from the shape of the input image. In order to address this issue, we can crop multiple rectangular areas with height and width that are integer multiples of 32 in the image, and perform forward propagation on the pixels in these areas separately. Note that the union of these rectangular areas needs to completely cover the input image. When a pixel is covered by multiple rectangular areas, the average of the transposed convolution outputs in separate areas for this same pixel can be input to the softmax operation to predict the class.

For simplicity, we only read a few larger test images, and crop a \(320\times480\) area for prediction starting from the upper-left corner of an image. For these test images, we print their cropped areas, prediction results, and ground-truth row by row.

voc_dir = d2l.download_extract('voc2012', 'VOCdevkit/VOC2012')
test_images, test_labels = d2l.read_voc_images(voc_dir, False)
n, imgs = 4, []
for i in range(n):
    crop_rect = (0, 0, 320, 480)
    X = torchvision.transforms.functional.crop(test_images[i], *crop_rect)
    pred = label2image(predict(X))
    imgs += [X.permute(1,2,0), pred.cpu(),
             torchvision.transforms.functional.crop(
                 test_labels[i], *crop_rect).permute(1,2,0)]
d2l.show_images(imgs[::3] + imgs[1::3] + imgs[2::3], 3, n, scale=2);

voc_dir = d2l.download_extract('voc2012', 'VOCdevkit/VOC2012')
test_images, test_labels = d2l.read_voc_images(voc_dir, False)
n, imgs = 4, []
for i in range(n):
    # Crop HWC arrays: top=0, left=0, height=320, width=480
    X = test_images[i][:320, :480, :]
    pred = label2image(predict(X))
    label_crop = test_labels[i][:320, :480, :]
    imgs += [X, pred.numpy(), label_crop]
d2l.show_images(imgs[::3] + imgs[1::3] + imgs[2::3], 3, n, scale=2);

voc_dir = d2l.download_extract('voc2012', 'VOCdevkit/VOC2012')
test_images, test_labels = d2l.read_voc_images(voc_dir, False)
n, imgs = 4, []
for i in range(n):
    # Crop HWC arrays: top=0, left=0, height=320, width=480
    X = test_images[i][:320, :480, :]
    pred = label2image(predict(X))
    label_crop = test_labels[i][:320, :480, :]
    imgs += [X, np.array(pred), label_crop]
d2l.show_images(imgs[::3] + imgs[1::3] + imgs[2::3], 3, n, scale=2);

voc_dir = d2l.download_extract('voc2012', 'VOCdevkit/VOC2012')
test_images, test_labels = d2l.read_voc_images(voc_dir, False)
n, imgs = 4, []
for i in range(n):
    crop_rect = (0, 0, 480, 320)
    X = image.fixed_crop(test_images[i], *crop_rect)
    pred = label2image(predict(X))
    imgs += [X, pred, image.fixed_crop(test_labels[i], *crop_rect)]
d2l.show_images(imgs[::3] + imgs[1::3] + imgs[2::3], 3, n, scale=2);

19.11.6 Summary

  • The fully convolutional network first uses a CNN to extract image features, then transforms the number of channels into the number of classes via a \(1\times 1\) convolutional layer, and finally transforms the height and width of the feature maps to those of the input image via the transposed convolution.
  • In a fully convolutional network, we can use upsampling of bilinear interpolation to initialize the transposed convolutional layer.

19.11.7 Exercises

  1. If we use Xavier initialization for the transposed convolutional layer in the experiment, how does the result change?
  2. Can you further improve the accuracy of the model by tuning the hyperparameters?
  3. Predict the classes of all pixels in test images.
  4. The original fully convolutional network paper also uses outputs of some intermediate CNN layers (Long et al. 2015). Try to implement this idea.