What is luz?

luz is a high-level training API built on top of torch. It removes most of the training-loop boilerplate we wrote in the previous lecture:

• Automatic batching via dataloaders

• Built-in metric tracking, validation splits, callbacks

• Early stopping, learning rate schedulers, checkpoints

• A built-in loss-curve plot

A quick recap

Recall the 10000-row regression problem we fit by hand last lecture using nn_linear and a manual training loop:

Xt$shape

[1] 10000    20

yt$shape

[1] 10000     1

coef(lm_fit) |> round(2)

(Intercept)          X1          X2          X3          X4 
       3.03        0.01        0.00        5.22       -0.01 
         X5          X6          X7          X8          X9 
       0.00        0.00       -3.10       -0.01        0.00 
        X10         X11         X12         X13         X14 
      -0.01       -0.01        8.69       -0.02       -0.01 
        X15         X16         X17         X18         X19 
       0.02        0.02        0.01       -1.48        0.00 
        X20 
      -0.01 

deviance(lm_fit) |> round(2)

[1] 9859.22

The luz workflow

fitted = nn_module(
  initialize = function(in_features) {
    self$linear = nn_linear(in_features, 1)
  },
  forward = function(X) {
    self$linear(X)
  }
) |>
  setup(
    loss = nn_mse_loss(),
    optimizer = optim_sgd
  ) |>
  set_hparams(in_features = p) |>
  set_opt_hparams(lr = 0.05) |>
  fit(
    list(Xt, yt),
    epochs = 20,
    dataloader_options = list(
      batch_size = 100,
      shuffle = TRUE
    ),
    verbose = FALSE
  )

Inspecting the fit

c(fitted$model$linear$bias, fitted$model$linear$weight) |> map(as.numeric) |> unlist() |> round(2)

 [1]  2.98  0.02  0.00  5.26 -0.04  0.02  0.04 -3.09  0.01  0.02 -0.04
[12] -0.02  8.70 -0.02  0.02  0.02  0.03  0.00 -1.54 -0.02 -0.01

lm_fit |> coef() |> round(2) |> unname()

 [1]  3.03  0.01  0.00  5.22 -0.01  0.00  0.00 -3.10 -0.01  0.00 -0.01
[12] -0.01  8.69 -0.02 -0.01  0.02  0.02  0.01 -1.48  0.00 -0.01

fitted$records$metrics$train |> unlist() |> unname()

 [1] 7.6132669 0.9980766 0.9995083 0.9977403 0.9978724 1.0000237
 [7] 0.9983734 0.9982529 0.9994720 0.9989458 0.9992331 0.9981256
[13] 0.9984254 0.9998939 1.0000370 0.9997186 0.9977539 0.9969737
[19] 1.0005277 0.9968545

mean(lm_fit$residuals^2)

[1] 0.9859224

Loss curve from luz

plot(fitted)

The digits dataset

We will now take a look at the MNIST dataset - this variant comes from Python sklearn package and is available as a CSV file in static/slides/data/mnist_digits.csv.

It contains 1797 8×8 grayscale images of handwritten digits.

digits = readr::read_csv("data/mnist_digits.csv", show_col_types = FALSE)
dim(digits)

[1] 1797   65

X = digits  |> select(-label) |> as.matrix()
y = digits$label
dim(X)

[1] 1797   64

table(y)

y
  0   1   2   3   4   5   6   7   8   9 
178 182 177 183 181 182 181 179 174 180 

Example digits

Train / test split

set.seed(1234)
split = initial_split(digits, prop = 0.8)

X_train = training(split) |> select(-label) |> as.matrix()
X_test  = testing(split)  |> select(-label) |> as.matrix()
y_train = training(split)$label
y_test  = testing(split)$label

dim(X_train); dim(X_test)

[1] 1437   64

[1] 360  64

To torch tensors

X_train_t = torch_tensor(X_train, dtype = torch_float())
y_train_t = torch_tensor(y_train + 1L, dtype = torch_long())
X_test_t  = torch_tensor(X_test,  dtype = torch_float())
y_test_t  = torch_tensor(y_test  + 1L, dtype = torch_long())

X_train_t$shape

[1] 1437   64

y_train_t$shape

[1] 1437

nn_linear

A quick refresher - nn_linear(in_features, out_features) implements the affine map

\[y = x W^\top + b\]

with a learnable weight matrix W of shape (out_features, in_features) and bias vector of length out_features. This is just a slightly funny way of writing a traditional regression model.

• Input shape: (N, in_features) - one row per example

• Output shape: (N, out_features) - one column per output unit

In our Regression model we used out_features = 1 since we were trying to predict a single number. For our classification model we will use out_features = 10 - one per class - and interpret the output as raw class scores (logits).

A linear (multinomial) model

The following is the torch analogue of nnet::multinom() - a single nn_linear layer from 64 pixel features to 10 class scores, trained with cross-entropy loss.

mnist_linear = nn_module(
  initialize = function() {
    self$linear = nn_linear(64, 10)
  },
  forward = function(x) {
    self$linear(x)
  }
)

Cross-entropy loss

For K-class classification the standard loss is cross-entropy. Given raw class scores (logits) \(z = (z_1, \dots, z_K)\) from the model, we convert to class probabilities with the softmax

\[p_k = \frac{\exp(z_k)}{\sum_{j=1}^{K} \exp(z_j)}\]

and the loss for a single example with true class \(y\) is the negative log-likelihood of the correct class

\[\ell(z, y) = -\log p_y = -z_y + \log \sum_{j=1}^{K} \exp(z_j).\]

Averaging over the mini-batch gives the objective that nn_cross_entropy_loss() minimizes.

nn_cross_entropy_loss()

A few things to keep in mind when using it:

• The model’s forward() should return raw logits, not probabilities - the softmax is fused into the loss for numerical stability.

• Targets are class indices (a long tensor), not a dummy coded / one-hot encoded matrix

• R torch uses 1-based indexing, so digit labels 0-9 become 1-10 (this is why we added 1 earlier).

loss_fn = nn_cross_entropy_loss()
logits = torch_tensor(matrix(c(2.0, 0.5, -1.0, 0.1), nrow = 1))
target = torch_tensor(1L)
loss_fn(logits, target)

torch_tensor
0.35240593552589417
[ CPUFloatType{} ]

-log(exp(2.0) / sum(exp(c(2.0, 0.5, -1.0, 0.1))))

[1] 0.3524059

Fitting with luz

torch_manual_seed(1)
fitted_lin = mnist_linear |>
  setup(
    loss = nn_cross_entropy_loss(),
    optimizer = optim_sgd,
    metrics = list(luz_metric_accuracy())
  ) |>
  set_opt_hparams(lr = 0.05) |>
  fit(
    list(X_train_t, y_train_t),
    valid_data = list(X_test_t, y_test_t),
    epochs = 30,
    dataloader_options = list(batch_size = 64, shuffle = TRUE),
    verbose = FALSE
  )

get_metrics(fitted_lin) |> filter(set == "train") |> tail(4)

     set metric epoch      value
57 train   loss    29 0.03515765
58 train    acc    29 0.99005682
59 train   loss    30 0.05474032
60 train    acc    30 0.97940341

get_metrics(fitted_lin) |> filter(set == "valid") |> tail(4)

     set metric epoch      value
57 valid   loss    29 0.06857039
58 valid    acc    29 0.98333333
59 valid   loss    30 0.08136037
60 valid    acc    30 0.97222222

Linear model results

Actication layers

We’re about to start building more complex models by stacking layers on top of each other. The issue is that if we just stack nn_linear layers, the whole thing collapses back down to a single linear transformation - no matter how many layers we add, the model is still just a linear model.

To solve this problem, we need to add a nonlinearity between the linear layers - an activation function that is applied element-wise to the output of each linear layer.

The two common choices:

• nn_relu() - Rectified Linear Unit: \(f(x) = \max(0, x)\). Fast, sparse, and the default in most modern networks.

• nn_tanh() - Hyperbolic tangent: \(f(x) = \tanh(x)\). Outputs in \((-1, 1)\), symmetric around zero.

Activations in practice

A feed-forward network

Now we’ll build a model flexible model by introducing a hidden layer with 64 units and a ReLU activation.

mnist_fnn = nn_module(
  initialize = function(hidden = 64, activation = nn_relu) {
    self$net = nn_sequential(
      nn_linear(64, hidden),
      activation(),
      nn_linear(hidden, 10)
    )
  },
  forward = function(x) {
    self$net(x)
  }
)

Fitting the FNN

torch_manual_seed(1)
fitted_fnn = mnist_fnn |>
  setup(
    loss = nn_cross_entropy_loss(),
    optimizer = optim_sgd,
    metrics = list(luz_metric_accuracy())
  ) |>
  set_hparams(hidden = 64, activation = nn_relu) |>
  set_opt_hparams(lr = 0.05) |>
  fit(
    list(X_train_t, y_train_t),
    valid_data = list(X_test_t, y_test_t),
    epochs = 50,
    dataloader_options = list(batch_size = 64, shuffle = TRUE),
    verbose = FALSE
  )

get_metrics(fitted_fnn) |> 
  filter(set == "train") |> tail(4)

      set metric epoch       value
97  train   loss    49 0.005963060
98  train    acc    49 1.000000000
99  train   loss    50 0.005669653
100 train    acc    50 1.000000000

get_metrics(fitted_fnn) |> 
  filter(set == "valid") |> tail(4)

      set metric epoch      value
97  valid   loss    49 0.05062127
98  valid    acc    49 0.98055556
99  valid   loss    50 0.05027586
100 valid    acc    50 0.97777778

The CIFAR10 dataset

• 60,000 color images at 32×32 resolution

• 10 classes - airplane, automobile, bird, cat, deer, dog, frog, horse, ship, truck

• 50,000 training / 10,000 test split

• A standard benchmark for evaluating CNN architectures

Smaller brother to the CIFAR100 dataset (100 classes) and the ImageNet dataset (1000 classes, 224×224 resolution) - the latter was the standard benchmark for large-scale image classification.

Loading the cached data

train = readRDS("data/cifar10/train.rds")
test  = readRDS("data/cifar10/test.rds")

str(train, max.level = 1)

List of 3
 $ x      : int [1:50000, 1:3, 1:32, 1:32] 59 154 255 28 170 159 164 28 134 125 ...
  ..- attr(*, "dimnames")=List of 4
 $ y      : int [1:50000] 7 10 10 5 2 2 3 8 9 4 ...
 $ classes: chr [1:10] "airplane" "automobile" "bird" "cat" ...

train$classes

 [1] "airplane"   "automobile" "bird"       "cat"        "deer"      
 [6] "dog"        "frog"       "horse"      "ship"       "truck"     

Example images

A custom dataset()

For datasets that don’t fit neatly into an (x, y) tuple we define a dataset() class with initialize, .getitem, and .length methods - the R torch analogue of Python’s torch.utils.data.Dataset:

cifar10_dataset = dataset(
  name = "cifar10",
  initialize = function(data) {
    self$x = torch_tensor(data$x, dtype = torch_float()) / 255
    self$y = torch_tensor(data$y, dtype = torch_long())
  },
  .getitem = function(i) {
    list(x = self$x[i, , , ], y = self$y[i])
  },
  .length = function() {
    self$y$size(1)
  }
)

From dataset to dataloader

train_ds = cifar10_dataset(train)
test_ds  = cifar10_dataset(test)

train_dl = dataloader(train_ds, batch_size = 128, shuffle = TRUE)
test_dl  = dataloader(test_ds,  batch_size = 128, shuffle = FALSE)

length(train_dl)

[1] 391

batch = train_dl$.iter()$.next()
batch$x$shape

[1] 128   3  32  32

batch$y$shape

[1] 128

A CIFAR10 feed-forward model

As a baseline, we flatten each 32×32×3 image into a 3072-length vector and feed it through a two-hidden-layer network - the same approach we used for MNIST digits.

cifar_fnn = nn_module(
  initialize = function() {
    self$net = nn_sequential(
      nn_flatten(),
      nn_linear(3 * 32 * 32, 512),
      nn_relu(),
      nn_linear(512, 128),
      nn_relu(),
      nn_linear(128, 10)
    )
  },
  forward = function(x) {
    self$net(x)
  }
)

2D convolutions

The model we just used flattened all of the data which throws away all spatial structure - the model has no way to know that nearby pixels are related. This limits how well a fully connected network can do on image data.

nn_conv2d

A 2D convolutional layer slides a bank of learnable filters over the spatial dimensions of its input - the same filter weights are reused at every location:

nn_conv2d(in_channels, out_channels, kernel_size, padding = 0, stride = 1)

• Input: (N, in_channels, H, W) - batch of multi-channel images

• Output: (N, out_channels, H', W') - one feature map per output channel

• Parameters: out_channels × (in_channels × kernel_size² + 1)

nn_conv2d in practice

conv = nn_conv2d(in_channels = 3, out_channels = 8,
                 kernel_size = 3, padding = 1)
conv$weight$shape

[1] 8 3 3 3

conv$bias$shape

[1] 8

x = torch_randn(5, 3, 32, 32)
conv(x)$shape

[1]  5  8 32 32

Pooling

After a convolutional layer we typically downsample the features. Pooling layers summarize small spatial neighborhoods into a single value - this has no learnable parameters.

Why pool?

• Reduces spatial dimensions, cutting computation and memory for later layers

• Gives a degree of translation invariance - small shifts in the input produce the same pooled output

• Forces the network to learn coarser, more abstract features as we go deeper

nn_max_pool2d and nn_avg_pool2d

The two common pooling operations over a 2×2 window:

• nn_max_pool2d(kernel_size = 2) - takes the maximum value in each window
• nn_avg_pool2d(kernel_size = 2) - takes the mean of values in each window

Both halve the spatial dimensions while leaving the channel count unchanged.

x = torch_randn(5, 8, 32, 32)

nn_max_pool2d(kernel_size = 2)(x)$shape

[1]  5  8 16 16

nn_avg_pool2d(kernel_size = 2)(x)$shape

[1]  5  8 16 16

A CIFAR10 CNN

Directly adapted from the PyTorch CIFAR10 tutorial - two conv+pool blocks feeding into a 3-layer mlp with ReLU activations.

cifar_cnn = nn_module(
  initialize = function() {
    self$net = nn_sequential(       # (N, 3, 32, 32)
      nn_conv2d(3, 6, kernel_size = 5),  # -> (N, 6, 28, 28)
      nn_relu(),
      nn_max_pool2d(kernel_size = 2, stride = 2),  # -> (N, 6, 14, 14)
      nn_conv2d(6, 16, kernel_size = 5), # -> (N, 16, 10, 10)
      nn_relu(),
      nn_max_pool2d(kernel_size = 2, stride = 2),  # -> (N, 16, 5, 5)
      nn_flatten(),                      # -> (N, 400)
      nn_linear(16 * 5 * 5, 120),       # -> (N, 120)
      nn_relu(),
      nn_linear(120, 84),               # -> (N, 84)
      nn_relu(),
      nn_linear(84, 10)                 # -> (N, 10)
    )
  },
  forward = function(x) {
    self$net(x)
  }
)

The VGG16 model

vgg16 = nn_module(
  initialize = function() {
    cfg = c(64, 64, "M", 128, 128, "M", 256, 256, 256, "M",
            512, 512, 512, "M", 512, 512, 512, "M")
    in_channels = 3
    layers = list()
    for (x in cfg) {
      if (x == "M") {
        layers[[length(layers) + 1]] = nn_max_pool2d(kernel_size = 2, stride = 2)
      } else {
        x = as.integer(x)
        layers[[length(layers) + 1]] = nn_conv2d(in_channels, x,
                                                 kernel_size = 3, padding = 1)
        layers[[length(layers) + 1]] = nn_batch_norm2d(x)
        layers[[length(layers) + 1]] = nn_relu(inplace = TRUE)
        in_channels = x
      }
    }
    layers[[length(layers) + 1]] = nn_flatten()
    layers[[length(layers) + 1]] = nn_linear(512, 10)
    self$net = do.call(nn_sequential, layers)
  },
  forward = function(x) {
    self$net(x)
  }
)

Fitting VGG16

On a laptops, each epoch of VGG16 on the full CIFAR10 training set takes several minutes on CPU. The code below would produce roughly the following loss / accuracy with lr = 0.01:

torch_manual_seed(1)
fitted_vgg = vgg16 |>
  setup(
    loss = nn_cross_entropy_loss(),
    optimizer = optim_sgd,
    metrics = list(luz_metric_accuracy())
  ) |>
  set_opt_hparams(lr = 0.01) |>
  fit(
    train_dl,
    valid_data = test_dl,
    epochs = 10
  )

Epoch  1 | train loss: 1.345 | valid loss: 1.102 | train acc: 0.512 | valid acc: 0.621
Epoch  2 | train loss: 0.790 | valid loss: 0.812 | train acc: 0.726 | valid acc: 0.714
Epoch  3 | train loss: 0.577 | valid loss: 0.731 | train acc: 0.800 | valid acc: 0.751
Epoch  4 | train loss: 0.445 | valid loss: 0.669 | train acc: 0.844 | valid acc: 0.781
Epoch  5 | train loss: 0.350 | valid loss: 0.701 | train acc: 0.878 | valid acc: 0.782
Epoch  6 | train loss: 0.274 | valid loss: 0.647 | train acc: 0.904 | valid acc: 0.803
Epoch  7 | train loss: 0.215 | valid loss: 0.668 | train acc: 0.925 | valid acc: 0.809
Epoch  8 | train loss: 0.167 | valid loss: 0.664 | train acc: 0.941 | valid acc: 0.822
Epoch  9 | train loss: 0.127 | valid loss: 0.710 | train acc: 0.955 | valid acc: 0.824
Epoch 10 | train loss: 0.103 | valid loss: 0.703 | train acc: 0.964 | valid acc: 0.832

Where to go next

• torch - the full low-level API. Everything we have done translates one-for-one with PyTorch documentation.

• luz - higher-level training, callbacks, early stopping, learning-rate schedulers, checkpointing.

• torchvision - image datasets, transforms, and a growing set of pretrained models (ResNet, VGG, EfficientNet, …).

• torchdatasets - a growing set of tabular and vision datasets with ready-made dataset() classes.

• GPU support - everything here runs unchanged on a CUDA-capable machine; tensors and modules are moved with $to(device = "cuda").

  • Departmental Servers have 2x NVidia A4000 GPUs each and you are welcome to make use of them for learning purposes