Cerca de duas semanas atrás, nós Probabilidade de Tensorflow introduzida (TFP)mostrando como criar e amostrar de distribuições e colocá -los para usar em um autoencoder variacional (VAE) que aprende seu anterior. Hoje, passamos para uma amostra diferente no zoológico do modelo VAE: o vetor quantizado autoencoder variacional (vq-vae) descrito em Aprendizado neural de representação discreta (Oord, Vinyals e Kavukcuoglu 2017). Esse modelo difere da maioria dos VAEs, pois seu posterior aproximado não é contínuo, mas discreto – daí o “quantizado” no título do artigo. Vamos analisar rapidamente o que isso significa e depois mergulharemos diretamente no código, combinando camadas de Keras, execução ansiosa e TFP.
Muitos fenômenos são melhor pensados e modelados, como discretos. Isso vale para fonemas e lexemes na linguagem, estruturas de nível superior nas imagens (pense em objetos em vez de pixels) e tarefas que exigem raciocínio e planejamento. O código latente usado na maioria dos Vaes, no entanto, é contínuo – geralmente é um gaussiano multivariado. Os vaes do espaço contínuo foram considerados muito bem-sucedidos em reconstruir sua entrada, mas muitas vezes eles sofrem de algo chamado colapso posterior: O decodificador é tão poderoso que pode criar uma saída realista, dada apenas qualquer entrada. Isso significa que não há incentivo para aprender um espaço latente expressivo.
Em VQ-VAE, no entanto, cada amostra de entrada é mapeada deterministicamente para um de um conjunto de incorporação de vetores. Juntos, esses vetores incorporados constituem o anterior para o espaço latente. Como tal, um vetor de incorporação contém muito mais informações do que uma média e uma variação e, portanto, é muito mais difícil de ignorar pelo decodificador.
A questão então é: onde está esse chapéu mágico, para que puxemos incorporações significativas?
A partir da descrição conceitual acima, agora temos duas perguntas a serem respondidas. Primeiro, por que mecanismo atribuímos amostras de entrada (que passaram pelo codificador) para a incorporação de vetores de incorporação apropriados? E segundo: como podemos aprender a incorporar vetores que realmente são representações úteis – que, quando alimentadas a um decodificador, resultarão em entidades percebidas como pertencentes à mesma espécie?
No que diz respeito à atribuição, um tensor emitido do codificador é simplesmente mapeado para o vizinho mais próximo na incorporação do espaço, usando a distância euclidiana. Os vetores de incorporação são atualizados usando médias móveis exponenciais. Como veremos em breve, isso significa que eles realmente não estão sendo aprendidos usando a ascendência de gradiente – um recurso que vale a pena apontar, pois não o encontramos todos os dias em aprendizado profundo.
Concretamente, como então a função de perda e o processo de treinamento deve ter? Provavelmente isso será mais fácil no código.
O código completo para este exemplo, incluindo utilitários para economia de modelo e visualização de imagem, é Disponível no GitHub Como parte dos exemplos de Keras. A ordem de apresentação aqui pode diferir da ordem de execução actual para fins expositivos; portanto, para executar o código, considere fazer uso do exemplo no GitHub.
Como em todas as nossas postagens anteriores no VAES, usamos a execução ansiosa, que pressupõe a implementação do TensorFlow das Keras.
Como em nosso publish anterior sobre o VAE com TFP, usaremos Kuzushiji-mnist(Clanuwat et al. 2018) como entrada. Agora é a hora de olhar para O que acabamos gerando esse tempo E faça sua aposta: como isso se compara ao espaço latente discreto do VQ-VAE?
np <- import("numpy")
kuzushiji <- np$load("kmnist-train-imgs.npz")
kuzushiji <- kuzushiji$get("arr_0")
train_images <- kuzushiji %>%
k_expand_dims() %>%
k_cast(dtype = "float32")
train_images <- train_images %>% `/`(255)
buffer_size <- 60000
batch_size <- 64
num_examples_to_generate <- batch_size
batches_per_epoch <- buffer_size / batch_size
train_dataset <- tensor_slices_dataset(train_images) %>%
dataset_shuffle(buffer_size) %>%
dataset_batch(batch_size, drop_remainder = TRUE)
Hyperparameters
Além dos hiperparâmetros “usuais” que temos na aprendizagem profunda, a infraestrutura do VQ-VAE apresenta alguns do modelo específico. Primeiro de tudo, o espaço de incorporação é de dimensionalidade Número de vetores de incorporação vezes Incorporar o tamanho do vetor:
# variety of embedding vectors
num_codes <- 64L
# dimensionality of the embedding vectors
code_size <- 16L
O espaço latente em nosso exemplo será do tamanho um, ou seja, temos um único vetor de incorporação representando o código latente para cada amostra de entrada. Isso ficará bem para o nosso conjunto de dados, mas deve -se notar que van den Oord et al. Utilizou espaços latentes muito mais altos no EG Imagenet e Cifar-10.
Modelo do codificador
O codificador usa camadas convolucionais para extrair recursos de imagem. Sua saída é um tensor 3D de forma BatchSize * 1 * code_size.
activation <- "elu"
# modularizing the code just a bit bit
default_conv <- set_defaults(layer_conv_2d, listing(padding = "similar", activation = activation))
base_depth <- 32
encoder_model <- perform(title = NULL,
code_size) {
keras_model_custom(title = title, perform(self) {
self$conv1 <- default_conv(filters = base_depth, kernel_size = 5)
self$conv2 <- default_conv(filters = base_depth, kernel_size = 5, strides = 2)
self$conv3 <- default_conv(filters = 2 * base_depth, kernel_size = 5)
self$conv4 <- default_conv(filters = 2 * base_depth, kernel_size = 5, strides = 2)
self$conv5 <- default_conv(filters = 4 * latent_size, kernel_size = 7, padding = "legitimate")
self$flatten <- layer_flatten()
self$dense <- layer_dense(items = latent_size * code_size)
self$reshape <- layer_reshape(target_shape = c(latent_size, code_size))
perform (x, masks = NULL) {
x %>%
# output form: 7 28 28 32
self$conv1() %>%
# output form: 7 14 14 32
self$conv2() %>%
# output form: 7 14 14 64
self$conv3() %>%
# output form: 7 7 7 64
self$conv4() %>%
# output form: 7 1 1 4
self$conv5() %>%
# output form: 7 4
self$flatten() %>%
# output form: 7 16
self$dense() %>%
# output form: 7 1 16
self$reshape()
}
})
}
Como sempre, vamos usar o fato de que estamos usando a execução ansiosa e ver algumas saídas de exemplo.
iter <- make_iterator_one_shot(train_dataset)
batch <- iterator_get_next(iter)
encoder <- encoder_model(code_size = code_size)
encoded <- encoder(batch)
encoded
tf.Tensor(
((( 0.00516277 -0.00746826 0.0268365 ... -0.012577 -0.07752544
-0.02947626))
...
((-0.04757921 -0.07282603 -0.06814402 ... -0.10861694 -0.01237121
0.11455103))), form=(64, 1, 16), dtype=float32)
Agora, cada um desses vetores 16D precisa ser mapeado para o vetor de incorporação do qual é mais próximo. Este mapeamento é resolvido por outro modelo: vector_quantizer
.
Modelo de quantizador vetorial
É assim que instanciaremos o quantizador vetorial:
vector_quantizer <- vector_quantizer_model(num_codes = num_codes, code_size = code_size)
Este modelo serve a dois propósitos: primeiro, atua como uma loja para os vetores de incorporação. Segundo, ele corresponde à saída do codificador às incorporações disponíveis.
Aqui, o estado atual das incorporações é armazenado em codebook
. ema_means
e ema_count
são apenas para fins de contabilidade (observe como eles estão definidos como não são transíveis). Vamos vê -los em uso em breve.
vector_quantizer_model <- perform(title = NULL, num_codes, code_size) {
keras_model_custom(title = title, perform(self) {
self$num_codes <- num_codes
self$code_size <- code_size
self$codebook <- tf$get_variable(
"codebook",
form = c(num_codes, code_size),
dtype = tf$float32
)
self$ema_count <- tf$get_variable(
title = "ema_count", form = c(num_codes),
initializer = tf$constant_initializer(0),
trainable = FALSE
)
self$ema_means = tf$get_variable(
title = "ema_means",
initializer = self$codebook$initialized_value(),
trainable = FALSE
)
perform (x, masks = NULL) {
# to be stuffed in shortly ...
}
})
}
Além das incorporações reais, em seu name
método vector_quantizer
segura a lógica de atribuição. Primeiro, calculamos a distância euclidiana de cada codificação para os vetores no livro de códigos (tf$norm
). Atribuímos cada codificação ao mais próximo, pois por essa distância incorporando (tf$argmin
) e um codificação de um scorching nas atribuições (tf$one_hot
). Finalmente, isolamos o vetor correspondente massando todos os outros e resumindo o que resta (multiplicação seguida de tf$reduce_sum
).
Em relação ao axis
argumento usado com muitas funções de tensorflow, leve em consideração que, em contraste com o seu k_*
irmãos, tensorflow cru (tf$*
) As funções esperam que a numeração do eixo seja baseada em 0. Também temos que adicionar o L
Após os números a estar em conformidade com os requisitos de tipo de dados da TensorFlow.
vector_quantizer_model <- perform(title = NULL, num_codes, code_size) {
keras_model_custom(title = title, perform(self) {
# right here we've got the above occasion fields
perform (x, masks = NULL) {
# form: bs * 1 * num_codes
distances <- tf$norm(
tf$expand_dims(x, axis = 2L) -
tf$reshape(self$codebook,
c(1L, 1L, self$num_codes, self$code_size)),
axis = 3L
)
# bs * 1
assignments <- tf$argmin(distances, axis = 2L)
# bs * 1 * num_codes
one_hot_assignments <- tf$one_hot(assignments, depth = self$num_codes)
# bs * 1 * code_size
nearest_codebook_entries <- tf$reduce_sum(
tf$expand_dims(
one_hot_assignments, -1L) *
tf$reshape(self$codebook, c(1L, 1L, self$num_codes, self$code_size)),
axis = 2L
)
listing(nearest_codebook_entries, one_hot_assignments)
}
})
}
Agora que vimos como os códigos são armazenados, vamos adicionar funcionalidade para atualizá -los. Como dissemos acima, eles não são aprendidos por descida de gradiente. Em vez disso, são médias móveis exponenciais, atualizadas continuamente por qualquer novo “membro da classe” que sejam atribuídas.
Então aqui está uma função update_ema
Isso vai cuidar disso.
update_ema
usa tensorflow Moving_averages para
- Primeiro, acompanhe o número de amostras atualmente atribuídas por código (
updated_ema_count
), e - segundo, calcule e atribua a média móvel exponencial atual (
updated_ema_means
).
moving_averages <- tf$python$coaching$moving_averages
# decay to make use of in computing exponential shifting common
decay <- 0.99
update_ema <- perform(
vector_quantizer,
one_hot_assignments,
codes,
decay) {
updated_ema_count <- moving_averages$assign_moving_average(
vector_quantizer$ema_count,
tf$reduce_sum(one_hot_assignments, axis = c(0L, 1L)),
decay,
zero_debias = FALSE
)
updated_ema_means <- moving_averages$assign_moving_average(
vector_quantizer$ema_means,
# selects all assigned values (masking out the others) and sums them up over the batch
# (can be divided by depend later, so we get a mean)
tf$reduce_sum(
tf$expand_dims(codes, 2L) *
tf$expand_dims(one_hot_assignments, 3L), axis = c(0L, 1L)),
decay,
zero_debias = FALSE
)
updated_ema_count <- updated_ema_count + 1e-5
updated_ema_means <- updated_ema_means / tf$expand_dims(updated_ema_count, axis = -1L)
tf$assign(vector_quantizer$codebook, updated_ema_means)
}
Antes de olharmos para o ciclo de treinamento, vamos concluir rapidamente a cena adicionando o último ator, o decodificador.
Modelo de decodificador
O decodificador é bastante padrão, realizando uma série de deconvoluções e, finalmente, retornando uma probabilidade para cada pixel de imagem.
default_deconv <- set_defaults(
layer_conv_2d_transpose,
listing(padding = "similar", activation = activation)
)
decoder_model <- perform(title = NULL,
input_size,
output_shape) {
keras_model_custom(title = title, perform(self) {
self$reshape1 <- layer_reshape(target_shape = c(1, 1, input_size))
self$deconv1 <-
default_deconv(
filters = 2 * base_depth,
kernel_size = 7,
padding = "legitimate"
)
self$deconv2 <-
default_deconv(filters = 2 * base_depth, kernel_size = 5)
self$deconv3 <-
default_deconv(
filters = 2 * base_depth,
kernel_size = 5,
strides = 2
)
self$deconv4 <-
default_deconv(filters = base_depth, kernel_size = 5)
self$deconv5 <-
default_deconv(filters = base_depth,
kernel_size = 5,
strides = 2)
self$deconv6 <-
default_deconv(filters = base_depth, kernel_size = 5)
self$conv1 <-
default_conv(filters = output_shape(3),
kernel_size = 5,
activation = "linear")
perform (x, masks = NULL) {
x <- x %>%
# output form: 7 1 1 16
self$reshape1() %>%
# output form: 7 7 7 64
self$deconv1() %>%
# output form: 7 7 7 64
self$deconv2() %>%
# output form: 7 14 14 64
self$deconv3() %>%
# output form: 7 14 14 32
self$deconv4() %>%
# output form: 7 28 28 32
self$deconv5() %>%
# output form: 7 28 28 32
self$deconv6() %>%
# output form: 7 28 28 1
self$conv1()
tfd$Impartial(tfd$Bernoulli(logits = x),
reinterpreted_batch_ndims = size(output_shape))
}
})
}
input_shape <- c(28, 28, 1)
decoder <- decoder_model(input_size = latent_size * code_size,
output_shape = input_shape)
Agora estamos prontos para treinar. Uma coisa que realmente não conversamos é a função de custo: dadas as diferenças na arquitetura (em comparação com os VAEs padrão), as perdas ainda parecerão o esperado (a adição common de perda de reconstrução e divergência de KL)? Vamos ver isso em um segundo.
Loop de treinamento
Aqui está o otimizador que usaremos. As perdas serão calculadas em linha.
optimizer <- tf$prepare$AdamOptimizer(learning_rate = learning_rate)
O loop de treinamento, como sempre, é um loop sobre as épocas, onde cada iteração é um loop sobre lotes obtidos no conjunto de dados. Para cada lote, temos um passe para a frente, gravado por um gradientTape
com base no qual calculamos a perda. A fita determinará os gradientes de todos os pesos treináveis em todo o modelo, e o otimizador usará esses gradientes para atualizar os pesos.
Até agora, tudo isso está em conformidade com um esquema que muitas vezes vimos antes. Um ponto a ser observado: nesse mesmo loop, também chamamos update_ema
Para recalcular as médias móveis, pois elas não são operadas durante o BackProp. Aqui está a funcionalidade essencial:
num_epochs <- 20
for (epoch in seq_len(num_epochs)) {
iter <- make_iterator_one_shot(train_dataset)
until_out_of_range({
x <- iterator_get_next(iter)
with(tf$GradientTape(persistent = TRUE) %as% tape, {
# do ahead go
# calculate losses
})
encoder_gradients <- tape$gradient(loss, encoder$variables)
decoder_gradients <- tape$gradient(loss, decoder$variables)
optimizer$apply_gradients(purrr::transpose(listing(
encoder_gradients, encoder$variables
)),
global_step = tf$prepare$get_or_create_global_step())
optimizer$apply_gradients(purrr::transpose(listing(
decoder_gradients, decoder$variables
)),
global_step = tf$prepare$get_or_create_global_step())
update_ema(vector_quantizer,
one_hot_assignments,
codes,
decay)
# periodically show some generated pictures
# see code on github
# visualize_images("kuzushiji", epoch, reconstructed_images, random_images)
})
}
Agora, para a ação actual. Dentro do contexto da fita gradiente, primeiro determinamos qual amostra de entrada codificada é atribuída a qual vetor de incorporação.
codes <- encoder(x)
c(nearest_codebook_entries, one_hot_assignments) %<-% vector_quantizer(codes)
Agora, para esta operação de atribuição, não há gradiente. Em vez disso, o que podemos fazer é passar os gradientes da entrada do decodificador diretamente até a saída do codificador. Aqui tf$stop_gradient
isenta nearest_codebook_entries
Da cadeia de gradientes, assim o codificador e o decodificador estão ligados por codes
:
codes_straight_through <- codes + tf$stop_gradient(nearest_codebook_entries - codes)
decoder_distribution <- decoder(codes_straight_through)
Em suma, o backprop cuidará dos pesos do decodificador e do codificador, enquanto as incorporações latentes são atualizadas usando médias móveis, como já vimos.
Agora estamos prontos para enfrentar as perdas. Existem três componentes:
- Primeiro, a perda de reconstrução, que é apenas a probabilidade de log da entrada actual sob a distribuição aprendida pelo decodificador.
reconstruction_loss <- -tf$reduce_mean(decoder_distribution$log_prob(x))
- Segundo, temos o perda de compromisso.
commitment_loss <- tf$reduce_mean(tf$sq.(codes - tf$stop_gradient(nearest_codebook_entries)))
- Finalmente, temos o KL recurring diverge para um anterior. Como a priori, todas as tarefas são igualmente prováveis, esse componente da perda é constante e muitas vezes pode ser dispensado. Estamos adicionando aqui principalmente para fins ilustrativos.
prior_dist <- tfd$Multinomial(
total_count = 1,
logits = tf$zeros(c(latent_size, num_codes))
)
prior_loss <- -tf$reduce_mean(
tf$reduce_sum(prior_dist$log_prob(one_hot_assignments), 1L)
)
Resumindo todos os três componentes, chegamos à perda geral:
beta <- 0.25
loss <- reconstruction_loss + beta * commitment_loss + prior_loss
Antes de olharmos para os resultados, vamos ver o que acontece dentro gradientTape
Em uma única olhada:
with(tf$GradientTape(persistent = TRUE) %as% tape, {
codes <- encoder(x)
c(nearest_codebook_entries, one_hot_assignments) %<-% vector_quantizer(codes)
codes_straight_through <- codes + tf$stop_gradient(nearest_codebook_entries - codes)
decoder_distribution <- decoder(codes_straight_through)
reconstruction_loss <- -tf$reduce_mean(decoder_distribution$log_prob(x))
commitment_loss <- tf$reduce_mean(tf$sq.(codes - tf$stop_gradient(nearest_codebook_entries)))
prior_dist <- tfd$Multinomial(
total_count = 1,
logits = tf$zeros(c(latent_size, num_codes))
)
prior_loss <- -tf$reduce_mean(tf$reduce_sum(prior_dist$log_prob(one_hot_assignments), 1L))
loss <- reconstruction_loss + beta * commitment_loss + prior_loss
})
Resultados
E aqui vamos nós. Desta vez, não podemos ter a 2D “Morphing View” que geralmente gosta de exibir com VAES (apenas não há espaço latente 2D). Em vez disso, as duas imagens abaixo são (1) letras geradas a partir de entrada aleatória e (2) reconstruídas actual Cartas, cada uma salva após o treinamento para nove épocas.

Duas coisas saltam para o olho: primeiro, as letras geradas são significativamente mais nítidas do que seus colegas de prior em prior (da postagem anterior). E segundo, você seria capaz de contar a imagem aleatória da imagem de reconstrução?
Nesse ponto, esperamos convencê-lo do poder e da eficácia dessa abordagem de latentes discretos. No entanto, você pode secretamente esperar que aplássemos isso a dados mais complexos, como os elementos da fala que mencionamos na introdução ou imagens de alta resolução, conforme encontrado no ImageNet.
A verdade é que há uma troca contínua entre o número de técnicas novas e emocionantes que podemos mostrar e o tempo que podemos gastar em iterações para aplicar com sucesso essas técnicas a conjuntos de dados complexos. No ultimate, é você, nossos leitores, que colocarão essas técnicas em uso significativo em dados relevantes do mundo actual.
Clanuwat, Tarin, Mikel Bober-Irizar, Asanobu Kitamoto, Alex Lamb, Kazuaki Yamamoto e David Ha. 2018. “Aprendizagem profunda para a literatura japonesa clássica”. 3 de dezembro de 2018. https://arxiv.org/abs/cs.cv/1812.01718.
Oord, Aaron Van den, Oriol Vinyals e Koray Kavukcuoglu. 2017. “Aprendizagem de representação discreta neural”. Corr ABS/1711.00937. http://arxiv.org/abs/1711.00937.