RStudio AI Weblog: Simple PixelCNN with tfprobability


We’ve seen fairly a couple of examples of unsupervised studying (or self-supervised studying, to decide on the extra appropriate however much less in style time period) on this weblog.

Usually, these concerned Variational Autoencoders (VAEs), whose attraction lies in them permitting to mannequin a latent area of underlying, impartial (ideally) components that decide the seen options. A attainable draw back may be the inferior high quality of generated samples. Generative Adversarial Networks (GANs) are one other in style strategy. Conceptually, these are extremely engaging because of their game-theoretic framing. Nonetheless, they are often tough to coach. PixelCNN variants, however – we’ll subsume all of them right here beneath PixelCNN – are typically identified for his or her good outcomes. They appear to contain some extra alchemy although. Beneath these circumstances, what might be extra welcome than a straightforward method of experimenting with them? Via TensorFlow Likelihood (TFP) and its R wrapper, tfprobability, we now have such a method.

This submit first provides an introduction to PixelCNN, concentrating on high-level ideas (leaving the small print for the curious to look them up within the respective papers). We’ll then present an instance of utilizing tfprobability to experiment with the TFP implementation.

PixelCNN rules

Autoregressivity, or: We’d like (some) order

The essential thought in PixelCNN is autoregressivity. Every pixel is modeled as relying on all prior pixels. Formally:

[p(mathbf{x}) = prod_{i}p(x_i|x_0, x_1, …, x_{i-1})]

Now wait a second – what even are prior pixels? Final I noticed one photographs have been two-dimensional. So this implies now we have to impose an order on the pixels. Generally this will probably be raster scan order: row after row, from left to proper. However when coping with coloration photographs, there’s one thing else: At every place, we even have three depth values, one for every of purple, inexperienced, and blue. The unique PixelCNN paper(Oord, Kalchbrenner, and Kavukcuoglu 2016) carried by means of autoregressivity right here as nicely, with a pixel’s depth for purple relying on simply prior pixels, these for inexperienced relying on these identical prior pixels however moreover, the present worth for purple, and people for blue relying on the prior pixels in addition to the present values for purple and inexperienced.

[p(x_i|mathbf{x}<i) = p(x_{i,R}|mathbf{x}<i) p(x_{i,G}|mathbf{x}<i, x_{i,R}) p(x_{i,B}|mathbf{x}<i, x_{i,R}, x_{i,G})]

Right here, the variant applied in TFP, PixelCNN++(Salimans et al. 2017) , introduces a simplification; it factorizes the joint distribution in a much less compute-intensive method.

Technically, then, we all know how autoregressivity is realized; intuitively, it might nonetheless appear shocking that imposing a raster scan order “simply works” (to me, at the least, it’s). Possibly that is a kind of factors the place compute energy efficiently compensates for lack of an equal of a cognitive prior.

Masking, or: The place to not look

Now, PixelCNN ends in “CNN” for a purpose – as typical in picture processing, convolutional layers (or blocks thereof) are concerned. However – is it not the very nature of a convolution that it computes a median of some kinds, wanting, for every output pixel, not simply on the corresponding enter but additionally, at its spatial (or temporal) environment? How does that rhyme with the look-at-just-prior-pixels technique?

Surprisingly, this downside is less complicated to resolve than it sounds. When making use of the convolutional kernel, simply multiply with a masks that zeroes out any “forbidden pixels” – like on this instance for a 5×5 kernel, the place we’re about to compute the convolved worth for row 3, column 3:

1 & 1 & 1 & 1 & 1
1 & 1 & 1 & 1 & 1
1 & 1 & 1 & 0 & 0
0 & 0 & 0 & 0 & 0
0 & 0 & 0 & 0 & 0

This makes the algorithm sincere, however introduces a distinct downside: With every successive convolutional layer consuming its predecessor’s output, there’s a constantly rising blind spot (so-called in analogy to the blind spot on the retina, however positioned within the prime proper) of pixels which are by no means seen by the algorithm. Van den Oord et al. (2016)(Oord et al. 2016) repair this through the use of two totally different convolutional stacks, one continuing from prime to backside, the opposite from left to proper.

Fig. 1: Left: Blind spot, growing over layers. Right: Using two different stacks (a vertical and a horizontal one) solves the problem. Source: van den Oord et al., 2016.

Conditioning, or: Present me a kitten

To date, we’ve all the time talked about “producing photographs” in a purely generic method. However the true attraction lies in creating samples of some specified sort – one of many lessons we’ve been coaching on, or orthogonal info fed into the community. That is the place PixelCNN turns into Conditional PixelCNN(Oord et al. 2016), and additionally it is the place that feeling of magic resurfaces. Once more, as “common math” it’s not onerous to conceive. Right here, (mathbf{h}) is the extra enter we’re conditioning on:

[p(mathbf{x}| mathbf{h}) = prod_{i}p(x_i|x_0, x_1, …, x_{i-1}, mathbf{h})]

However how does this translate into neural community operations? It’s simply one other matrix multiplication ((V^T mathbf{h})) added to the convolutional outputs ((W mathbf{x})).

[mathbf{y} = tanh(W_{k,f} mathbf{x} + V^T_{k,f} mathbf{h}) odot sigma(W_{k,g} mathbf{x} + V^T_{k,g} mathbf{h})]

(Should you’re questioning in regards to the second half on the proper, after the Hadamard product signal – we received’t go into particulars, however in a nutshell, it’s one other modification launched by (Oord et al. 2016), a switch of the “gating” precept from recurrent neural networks, resembling GRUs and LSTMs, to the convolutional setting.)

So we see what goes into the choice of a pixel worth to pattern. However how is that call truly made?

Logistic combination probability , or: No pixel is an island

Once more, that is the place the TFP implementation doesn’t comply with the unique paper, however the latter PixelCNN++ one. Initially, pixels have been modeled as discrete values, selected by a softmax over 256 (0-255) attainable values. (That this truly labored looks as if one other occasion of deep studying magic. Think about: On this mannequin, 254 is as removed from 255 as it’s from 0.)

In distinction, PixelCNN++ assumes an underlying steady distribution of coloration depth, and rounds to the closest integer. That underlying distribution is a mix of logistic distributions, thus permitting for multimodality:

[nu sim sum_{i} pi_i logistic(mu_i, sigma_i)]

General structure and the PixelCNN distribution

General, PixelCNN++, as described in (Salimans et al. 2017), consists of six blocks. The blocks collectively make up a UNet-like construction, successively downsizing the enter after which, upsampling once more:

Fig. 2: Overall structure of PixelCNN++. From: Salimans et al., 2017.

In TFP’s PixelCNN distribution, the variety of blocks is configurable as num_hierarchies, the default being 3.

Every block consists of a customizable variety of layers, known as ResNet layers because of the residual connection (seen on the proper) complementing the convolutional operations within the horizontal stack:

Fig. 3: One so-called "ResNet layer", featuring both a vertical and a horizontal convolutional stack. Source: van den Oord et al., 2017.

In TFP, the variety of these layers per block is configurable as num_resnet.

num_resnet and num_hierarchies are the parameters you’re almost definitely to experiment with, however there are a couple of extra you’ll be able to take a look at within the documentation. The variety of logistic distributions within the combination can be configurable, however from my experiments it’s greatest to maintain that quantity fairly low to keep away from producing NaNs throughout coaching.

Let’s now see an entire instance.

Finish-to-end instance

Our playground will probably be QuickDraw, a dataset – nonetheless rising – obtained by asking individuals to attract some object in at most twenty seconds, utilizing the mouse. (To see for your self, simply take a look at the web site). As of at present, there are greater than a fifty million situations, from 345 totally different lessons.

Initially, these knowledge have been chosen to take a break from MNIST and its variants. However similar to these (and lots of extra!), QuickDraw may be obtained, in tfdatasets-ready kind, by way of tfds, the R wrapper to TensorFlow datasets. In distinction to the MNIST “household” although, the “actual samples” are themselves extremely irregular, and infrequently even lacking important elements. So to anchor judgment, when displaying generated samples we all the time present eight precise drawings with them.

Making ready the information

The dataset being gigantic, we instruct tfds to load the primary 500,000 drawings “solely.”

To hurry up coaching additional, we then zoom in on twenty lessons. This successfully leaves us with ~ 1,100 – 1,500 drawings per class.

# bee, bicycle, broccoli, butterfly, cactus,
# frog, guitar, lightning, penguin, pizza,
# rollerskates, sea turtle, sheep, snowflake, solar,
# swan, The Eiffel Tower, tractor, practice, tree
lessons <- c(26, 29, 43, 49, 50,
             125, 134, 172, 218, 225,
             246, 255, 258, 271, 295,
             296, 308, 320, 322, 323

classes_tensor <- tf$solid(lessons, tf$int64)

train_ds <- train_ds %>%
    operate(report) tf$reduce_any(tf$equal(classes_tensor, report$label), -1L)

The PixelCNN distribution expects values within the vary from 0 to 255 – no normalization required. Preprocessing then consists of simply casting pixels and labels every to float:

preprocess <- operate(report) {
  report$picture <- tf$solid(report$picture, tf$float32) 
  report$label <- tf$solid(report$label, tf$float32)
  listing(tuple(report$picture, report$label))

batch_size <- 32

practice <- train_ds %>%
  dataset_map(preprocess) %>%
  dataset_shuffle(10000) %>%

Creating the mannequin

We now use tfd_pixel_cnn to outline what would be the loglikelihood utilized by the mannequin.

dist <- tfd_pixel_cnn(
  image_shape = c(28, 28, 1),
  conditional_shape = listing(),
  num_resnet = 5,
  num_hierarchies = 3,
  num_filters = 128,
  num_logistic_mix = 5,
  dropout_p =.5

image_input <- layer_input(form = c(28, 28, 1))
label_input <- layer_input(form = listing())
log_prob <- dist %>% tfd_log_prob(image_input, conditional_input = label_input)

This tradition loglikelihood is added as a loss to the mannequin, after which, the mannequin is compiled with simply an optimizer specification solely. Throughout coaching, loss first decreased shortly, however enhancements from later epochs have been smaller.

mannequin <- keras_model(inputs = listing(image_input, label_input), outputs = log_prob)
mannequin$compile(optimizer = optimizer_adam(lr = .001))

mannequin %>% match(practice, epochs = 10)

To collectively show actual and faux photographs:

for (i in lessons) {
  real_images <- train_ds %>%
      operate(report) report$label == tf$solid(i, tf$int64)
    ) %>% 
    dataset_take(8) %>%
  it <- as_iterator(real_images)
  real_images <- iter_next(it)
  real_images <- real_images$picture %>% as.array()
  real_images <- real_images[ , , , 1]/255
  generated_images <- dist %>% tfd_sample(8, conditional_input = i)
  generated_images <- generated_images %>% as.array()
  generated_images <- generated_images[ , , , 1]/255
  photographs <- abind::abind(real_images, generated_images, alongside = 1)
  png(paste0("draw_", i, ".png"), width = 8 * 28 * 10, peak = 2 * 28 * 10)
  par(mfrow = c(2, 8), mar = c(0, 0, 0, 0))
  photographs %>%
    purrr::array_tree(1) %>%
    purrr::map(as.raster) %>%

From our twenty lessons, right here’s a selection of six, every displaying actual drawings within the prime row, and faux ones under.

Fig. 4: Bicycles, drawn by people (top row) and the network (bottom row).
Fig. 5: Broccoli, drawn by people (top row) and the network (bottom row).
Fig. 6: Butterflies, drawn by people (top row) and the network (bottom row).
Fig. 7: Guitars, drawn by people (top row) and the network (bottom row).
Fig. 8: Penguins, drawn by people (top row) and the network (bottom row).
Fig. 9: Roller skates, drawn by people (top row) and the network (bottom row).

We in all probability wouldn’t confuse the primary and second rows, however then, the precise human drawings exhibit monumental variation, too. And nobody ever mentioned PixelCNN was an structure for idea studying. Be happy to mess around with different datasets of your selection – TFP’s PixelCNN distribution makes it simple.

Wrapping up

On this submit, we had tfprobability / TFP do all of the heavy lifting for us, and so, may give attention to the underlying ideas. Relying in your inclinations, this may be a perfect scenario – you don’t lose sight of the forest for the bushes. Alternatively: Do you have to discover that altering the offered parameters doesn’t obtain what you need, you might have a reference implementation to start out from. So regardless of the end result, the addition of such higher-level performance to TFP is a win for the customers. (Should you’re a TFP developer studying this: Sure, we’d like extra :-)).

To everybody although, thanks for studying!

Oord, Aaron van den, Nal Kalchbrenner, and Koray Kavukcuoglu. 2016. “Pixel Recurrent Neural Networks.” CoRR abs/1601.06759.
Oord, Aaron van den, Nal Kalchbrenner, Oriol Vinyals, Lasse Espeholt, Alex Graves, and Koray Kavukcuoglu. 2016. “Conditional Picture Technology with PixelCNN Decoders.” CoRR abs/1606.05328.

Salimans, Tim, Andrej Karpathy, Xi Chen, and Diederik P. Kingma. 2017. “PixelCNN++: A PixelCNN Implementation with Discretized Logistic Combination Chance and Different Modifications.” In ICLR.


Leave a Reply

Your email address will not be published. Required fields are marked *