Neural Style Transfer com Torch

Conceitos e ideias importantes

A ideia é gerar uma imagem generated G que tenha conteúdo similar a content C e estilo similar a style S. O conteúdo e o estilo são, geralmente, extraídos de uma convnet pré-treinada. O artigo original (Gatys, Ecker, and Bethge 2015) utiliza o VGG19, que já vem dentro do {torchvision}.

# VGG19 model
vgg <- model_vgg19(pretrained = TRUE)$features$to(device = device)

# congelando os pesos
vgg$parameters %>% purrr::walk(function(param) param$requires_grad_(FALSE))

Então o VGG19 vai nos fornecer as features e agora precisamos definir funções de custo para achar a imagem que tem o melhor compromisso entre o conteúdo de uma imagem e o estilo de outra. A construção do algoritmo pode ser divida em três partes:

• Função de custo do content: \(L_{content}(C, G)\)
• Função de custo do style: \(L_{style}(S, G)\)
• Função de custo global: \(L(G) = \alpha L_{content}(C, G) + \beta L_{style}(S, G)\)

content_loss <- function(content_layer, generated_layer) {
  nnf_mse_loss(content_layer, generated_layer)
}

gram_matrix <- function(tensor) {
  c(b,c,h,w) %<-% tensor$size()
  tensor <- tensor$view(c(c, h*w))
  torch_matmul(tensor, tensor$t())/(h*w)
}

style_loss <- function(style_layer, generated_layer) {
  style_gram <- gram_matrix(style_layer)
  generated_gram <- gram_matrix(generated_layer)
  nnf_mse_loss(style_gram, generated_gram)
}

style_layers <- c(2, 7, 12, 21, 29) # escolha das layers da VGG
lambdas <- 1e5/(c(1,10,10,10,10)^2) # pesos para cada uma das layers no estilo final

features <- nn_module(
  initialize = function(convnet) {
    # poderia ser qualquer convnet pré-treinada. Iremos usar a VGG19
    self$convnet <- convnet
  },
  forward = function(img, layers_out = NULL) {
    layers <- seq_along(self$convnet) # 1 a 37
    if(is.null(layers_out)) layers_out <- layers
    conv_outs <- purrr::accumulate(layers, ~self$convnet[[.y]](.x), .init = img) # lista de 37 tensores
    conv_outs[layers_out] # lista apenas com as layers selecionadas
  }
)

Otimização

Abaixo segue código comentado para rodar a otimização.

#funções auxiliares
to_r <- function(x) as.numeric(x$to(device = cpu))
                               
plot_image <- function(tensor) {
  im <- tensor$to(device = "cpu")[1,..]$
    permute(c(2, 3, 1))$
    to(device = "cpu")$
    clamp(0,1) %>% # make it [0,1]
    as.array()
  par(mar = c(0,0,0,0))
  plot(as.raster(im))
}

load_image <- function(path, geometry = "250x200") {
  img <- path %>%
    magick_loader() %>%
    magick::image_resize(geometry) %>%
    transform_to_tensor() %>%
    torch_unsqueeze(1)
  
  img$to(device = device)
}

Os parâmetros e inputs que valem a pena experimentar são:

• content e style: Imagens de input: a de conteúdo e a de estilo.
• content_layer: a layer da VGG19 para extrair o conteúdo imagem de conteúdo.
• style_layers: as layers da VGG19 para extrair o estilo imagem de estilo.
• lambdas: os pesos de cada feature de estilo na otimização.
• content_weight: o peso das features de conteúdo na otimização.
• style_weight: o peso das features de estilo (global) na otimização.

Outros dois parâmetros que afetam drasticamente o resultado são as dimensões das duas imagens de input. No código abaixo onde tem "400x400" e "350x500" pode-se trocar por outras dimensões a fim de se obter resultados diferentes.

No processo de criação você irá mexer nesses parâmetros o tempo todo!

# INPUT: content and style images
content <- load_image("content/posts/2021-02-22-neural-style-transfer/cristoredentor3.jpg", "400x400")
style <- load_image("content/posts/2021-02-22-neural-style-transfer/vangogh_starry_night.jpg", "350x500")

# style and content feature setup
content_layer <- 14
style_layers <- c(2, 7, 12, 21, 29)
lambdas <- 1e5/(c(1,10,10,10,10)^2)
content_weight <- 2
style_weight <- 4e-1

content

style

# FEATURES: extraídas da VGG19
vgg_features <- features(vgg)
content_features <- vgg_features(content, content_layer)
style_features <- vgg_features(style, style_layers)

# OUTPUT: generated image
generated <- torch_clone(content)$requires_grad_(TRUE)
optim <- optim_adam(generated, lr = 0.02)
lr_scheduler <- lr_step(optim, 100, 0.9)

# loop de otimização
for(step in seq_len(8000)) {
  optim$zero_grad()
  # atualiza as features da imagem que está sendo gerada
  generated_features <- vgg_features(generated, c(content_layer, style_layers))
  
  # losses
  LC <- content_loss(content_features[[1]], generated_features[[1]])
  LS <- 0
  for(i in seq_along(lambdas)) 
    LS <- LS + lambdas[i]*style_loss(style_features[[i]], generated_features[-1][[i]])  
  
  loss <- content_weight * LC + style_weight * LS
  
  loss$backward()
  optim$step()
  lr_scheduler$step()
  
  # feedback
  if(step %% 100 == 0) {
    cat(glue::glue("LC = {to_r(LC)} - LS = {to_r(LS)} - Loss = {to_r(loss)}\n\n"))
    plot_image(generated)
  }
}

# imagem final
plot_image(generated)

LC = 2.70741701126099 - LS = 36.7998428344727 - Loss = 20.1347713470459
LC = 2.61282753944397 - LS = 24.3190288543701 - Loss = 14.9532661437988
LC = 2.54735898971558 - LS = 19.7909469604492 - Loss = 13.0110969543457
LC = 2.49572896957397 - LS = 16.9353790283203 - Loss = 11.7656097412109
LC = 2.44669985771179 - LS = 16.6177845001221 - Loss = 11.5405139923096

E aí? Ficou com cara de pintura? Comente o que achou! Tente com suas imagens e compartilhe com a gente =). Aprender como NST funciona é um grande exercício para aprimorar o entendimento sobre modelos de deep learning em geral, principalmente sobre como podemos criar funções de custo mais especializadas em um determinado contexto.

Aprenda Deep learning com a Curso-R

Se você quiser entrar no incrível mundo das redes profundas, nosso curso de Deep Learning é a melhor porta de entrada, conheça!