Posit AI Weblog: Implementing rotation equivariance: Group-equivariant CNN from scratch

g2 <- elems[3]

# in C_4$left_action_on_H(), H stands for the symmetry group
C_4$left_action_on_H(torch_tensor(g1)$unsqueeze(1), torch_tensor(g2)$unsqueeze(1))

torch_tensor
 4.7124
[ CPUFloatType{1,1} ]

What’s with the unsqueeze()s? Since (C_4)’s final raison d’être is to be a part of a neural community, left_action_on_H() works with batches of components, not scalar tensors.

Issues are a bit much less simple the place the group motion on (mathbb{R}^2) is anxious. Right here, we want the idea of a group illustration. That is an concerned subject, which we received’t go into right here. In our present context, it really works about like this: We’ve an enter sign, a tensor we’d wish to function on ultimately. (That “a way” shall be convolution, as we’ll see quickly.) To render that operation group-equivariant, we first have the illustration apply the inverse group motion to the enter. That completed, we go on with the operation as if nothing had occurred.

To offer a concrete instance, let’s say the operation is a measurement. Think about a runner, standing on the foot of some mountain path, able to run up the climb. We’d wish to report their top. One choice we’ve got is to take the measurement, then allow them to run up. Our measurement shall be as legitimate up the mountain because it was down right here. Alternatively, we is perhaps well mannered and never make them wait. As soon as they’re up there, we ask them to come back down, and after they’re again, we measure their top. The consequence is similar: Physique top is equivariant (greater than that: invariant, even) to the motion of operating up or down. (In fact, top is a fairly uninteresting measure. However one thing extra attention-grabbing, corresponding to coronary heart fee, wouldn’t have labored so effectively on this instance.)

Returning to the implementation, it seems that group actions are encoded as matrices. There’s one matrix for every group aspect. For (C_4), the so-called customary illustration is a rotation matrix:

[
begin{bmatrix} cos(theta) & -sin(theta) sin(theta) & cos(theta) end{bmatrix}
]

In gcnn, the perform making use of that matrix is left_action_on_R2(). Like its sibling, it’s designed to work with batches (of group components in addition to (mathbb{R}^2) vectors). Technically, what it does is rotate the grid the picture is outlined on, after which, re-sample the picture. To make this extra concrete, that technique’s code seems to be about as follows.

Here’s a goat.

img_path <- system.file("imgs", "z.jpg", package deal = "gcnn")
img <- torchvision::base_loader(img_path) |> torchvision::transform_to_tensor()
img$permute(c(2, 3, 1)) |> as.array() |> as.raster() |> plot()

First, we name C_4$left_action_on_R2() to rotate the grid.

# Grid form is [2, 1024, 1024], for a second, 1024 x 1024 picture.
img_grid_R2 <- torch::torch_stack(torch::torch_meshgrid(
    listing(
      torch::torch_linspace(-1, 1, dim(img)[2]),
      torch::torch_linspace(-1, 1, dim(img)[3])
    )
))

# Remodel the picture grid with the matrix illustration of some group aspect.
transformed_grid <- C_4$left_action_on_R2(C_4$inverse(g1)$unsqueeze(1), img_grid_R2)

Second, we re-sample the picture on the remodeled grid. The goat now seems to be as much as the sky.

transformed_img <- torch::nnf_grid_sample(
  img$unsqueeze(1), transformed_grid,
  align_corners = TRUE, mode = "bilinear", padding_mode = "zeros"
)

transformed_img[1,..]$permute(c(2, 3, 1)) |> as.array() |> as.raster() |> plot()

Step 2: The lifting convolution

We wish to make use of current, environment friendly torch performance as a lot as attainable. Concretely, we wish to use nn_conv2d(). What we want, although, is a convolution kernel that’s equivariant not simply to translation, but in addition to the motion of (C_4). This may be achieved by having one kernel for every attainable rotation.

Implementing that concept is strictly what LiftingConvolution does. The precept is similar as earlier than: First, the grid is rotated, after which, the kernel (weight matrix) is re-sampled to the remodeled grid.

Why, although, name this a lifting convolution? The same old convolution kernel operates on (mathbb{R}^2); whereas our prolonged model operates on combos of (mathbb{R}^2) and (C_4). In math communicate, it has been lifted to the semi-direct product (mathbb{R}^2rtimes C_4).

lifting_conv <- LiftingConvolution(
    group = CyclicGroup(order = 4),
    kernel_size = 5,
    in_channels = 3,
    out_channels = 8
  )

x <- torch::torch_randn(c(2, 3, 32, 32))
y <- lifting_conv(x)
y$form

[1]  2  8  4 28 28

Since, internally, LiftingConvolution makes use of a further dimension to appreciate the product of translations and rotations, the output just isn’t four-, however five-dimensional.

Step 3: Group convolutions

Now that we’re in “group-extended house”, we will chain numerous layers the place each enter and output are group convolution layers. For instance:

group_conv <- GroupConvolution(
  group = CyclicGroup(order = 4),
    kernel_size = 5,
    in_channels = 8,
    out_channels = 16
)

z <- group_conv(y)
z$form

[1]  2 16  4 24 24

All that continues to be to be finished is package deal this up. That’s what gcnn::GroupEquivariantCNN() does.

Step 4: Group-equivariant CNN

We are able to name GroupEquivariantCNN() like so.

cnn <- GroupEquivariantCNN(
    group = CyclicGroup(order = 4),
    kernel_size = 5,
    in_channels = 1,
    out_channels = 1,
    num_hidden = 2, # variety of group convolutions
    hidden_channels = 16 # variety of channels per group conv layer
)

img <- torch::torch_randn(c(4, 1, 32, 32))
cnn(img)$form

[1] 4 1

At informal look, this GroupEquivariantCNN seems to be like several previous CNN … weren’t it for the group argument.

Now, once we examine its output, we see that the extra dimension is gone. That’s as a result of after a sequence of group-to-group convolution layers, the module initiatives all the way down to a illustration that, for every batch merchandise, retains channels solely. It thus averages not simply over places – as we usually do – however over the group dimension as effectively. A remaining linear layer will then present the requested classifier output (of dimension out_channels).

And there we’ve got the whole structure. It’s time for a real-world(ish) take a look at.

Rotated digits!

The concept is to coach two convnets, a “regular” CNN and a group-equivariant one, on the same old MNIST coaching set. Then, each are evaluated on an augmented take a look at set the place every picture is randomly rotated by a steady rotation between 0 and 360 levels. We don’t count on GroupEquivariantCNN to be “good” – not if we equip with (C_4) as a symmetry group. Strictly, with (C_4), equivariance extends over 4 positions solely. However we do hope it would carry out considerably higher than the shift-equivariant-only customary structure.

First, we put together the info; particularly, the augmented take a look at set.

dir <- "/tmp/mnist"

train_ds <- torchvision::mnist_dataset(
  dir,
  obtain = TRUE,
  remodel = torchvision::transform_to_tensor
)

test_ds <- torchvision::mnist_dataset(
  dir,
  practice = FALSE,
  remodel = perform(x) >
      torchvision::transform_random_rotation(
        levels = c(0, 360),
        resample = 2,
        fill = 0
      )
  
)

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

How does it look?

test_images <- coro::acquire(
  test_dl, 1
)[[1]]$x[1:32, 1, , ] |> as.array()

par(mfrow = c(4, 8), mar = rep(0, 4), mai = rep(0, 4))
test_images |>
  purrr::array_tree(1) |>
  purrr::map(as.raster) |>
  purrr::iwalk(~ {
    plot(.x)
  })

We first outline and practice a standard CNN. It’s as just like GroupEquivariantCNN(), architecture-wise, as attainable, and is given twice the variety of hidden channels, in order to have comparable capability general.

 default_cnn <- nn_module(
   "default_cnn",
   initialize = perform(kernel_size, in_channels, out_channels, num_hidden, hidden_channels) {
     self$conv1 <- torch::nn_conv2d(in_channels, hidden_channels, kernel_size)
     self$convs <- torch::nn_module_list()
     for (i in 1:num_hidden) {
       self$convs$append(torch::nn_conv2d(hidden_channels, hidden_channels, kernel_size))
     }
     self$avg_pool <- torch::nn_adaptive_avg_pool2d(1)
     self$final_linear <- torch::nn_linear(hidden_channels, out_channels)
   },
   ahead = perform(x) >
       torch::nnf_relu()
     for (i in 1:(size(self$convs))) >
         self$convs[[i]]() 
     x <- x 
 )

fitted <- default_cnn |>
    luz::setup(
      loss = torch::nn_cross_entropy_loss(),
      optimizer = torch::optim_adam,
      metrics = listing(
        luz::luz_metric_accuracy()
      )
    ) |>
    luz::set_hparams(
      kernel_size = 5,
      in_channels = 1,
      out_channels = 10,
      num_hidden = 4,
      hidden_channels = 32
    ) %>%
    luz::set_opt_hparams(lr = 1e-2, weight_decay = 1e-4) |>
    luz::match(train_dl, epochs = 10, valid_data = test_dl) 

Practice metrics: Loss: 0.0498 - Acc: 0.9843
Legitimate metrics: Loss: 3.2445 - Acc: 0.4479

Unsurprisingly, accuracy on the take a look at set just isn’t that nice.

Subsequent, we practice the group-equivariant model.

fitted <- GroupEquivariantCNN |>
  luz::setup(
    loss = torch::nn_cross_entropy_loss(),
    optimizer = torch::optim_adam,
    metrics = listing(
      luz::luz_metric_accuracy()
    )
  ) |>
  luz::set_hparams(
    group = CyclicGroup(order = 4),
    kernel_size = 5,
    in_channels = 1,
    out_channels = 10,
    num_hidden = 4,
    hidden_channels = 16
  ) |>
  luz::set_opt_hparams(lr = 1e-2, weight_decay = 1e-4) |>
  luz::match(train_dl, epochs = 10, valid_data = test_dl)

Practice metrics: Loss: 0.1102 - Acc: 0.9667
Legitimate metrics: Loss: 0.4969 - Acc: 0.8549

For the group-equivariant CNN, accuracies on take a look at and coaching units are lots nearer. That could be a good consequence! Let’s wrap up at the moment’s exploit resuming a thought from the primary, extra high-level publish.

A problem

Going again to the augmented take a look at set, or fairly, the samples of digits displayed, we discover an issue. In row two, column 4, there’s a digit that “below regular circumstances”, needs to be a 9, however, likely, is an upside-down 6. (To a human, what suggests that is the squiggle-like factor that appears to be discovered extra typically with sixes than with nines.) Nonetheless, you may ask: does this have to be an issue? Perhaps the community simply must be taught the subtleties, the sorts of issues a human would spot?

The best way I view it, all of it is dependent upon the context: What actually needs to be completed, and the way an utility goes for use. With digits on a letter, I’d see no purpose why a single digit ought to seem upside-down; accordingly, full rotation equivariance can be counter-productive. In a nutshell, we arrive on the identical canonical crucial advocates of honest, simply machine studying hold reminding us of:

All the time consider the best way an utility goes for use!

In our case, although, there may be one other side to this, a technical one. gcnn::GroupEquivariantCNN() is a straightforward wrapper, in that its layers all make use of the identical symmetry group. In precept, there isn’t any want to do that. With extra coding effort, completely different teams can be utilized relying on a layer’s place within the feature-detection hierarchy.

Right here, let me lastly inform you why I selected the goat image. The goat is seen by a red-and-white fence, a sample – barely rotated, as a result of viewing angle – made up of squares (or edges, if you happen to like). Now, for such a fence, kinds of rotation equivariance corresponding to that encoded by (C_4) make a variety of sense. The goat itself, although, we’d fairly not have look as much as the sky, the best way I illustrated (C_4) motion earlier than. Thus, what we’d do in a real-world image-classification process is use fairly versatile layers on the backside, and more and more restrained layers on the prime of the hierarchy.

Thanks for studying!

Picture by Marjan Blan | @marjanblan on Unsplash