-
Notifications
You must be signed in to change notification settings - Fork 13
/
train_MixModel_VAE.jl
69 lines (50 loc) · 1.87 KB
/
train_MixModel_VAE.jl
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
using MNIST
include("utils/optimizers.jl")
include("models/gaussMMVAE.jl")
function elbo(params, x, prior)
# prop through autoencoder
x_recon, posterior, z = fprop(params, x)
# likelihood term
expected_nll = sum(-x .* log(x_recon) - (1-x) .* log(1-x_recon), 2)
# kl divergence term
kld = computeKLD(prior, posterior, z)
return mean(expected_nll + kld)
end
function trainVAE(data, params, hyperParams)
N,d = size(data)
nBatches = div(N,hyperParams["batchSize"])
for epoch_idx in 1:hyperParams["nEpochs"]
elbo_tracker = 0.
for batch_idx in 1:nBatches
# get minibatch
x = data[(batch_idx-1)*hyperParams["batchSize"]+1:batch_idx*hyperParams["batchSize"],:]
# compute elbo
elbo_tracker += elbo(params, x, hyperParams["prior"])
# get elbo gradients
#grads = elbo_grad(params, x, hyperParams["prior"])
# perform AdaM update
#params = adamUpdate(params, grads, hyperParams["adamParams"])
end
@printf "Epoch %d. Neg. ELBO: %.3f \n" epoch_idx elbo_tracker/nBatches
end
return params
end
function main()
# load MNIST (values in [0,255])
data = transpose(traindata()[1])
# shuffle, normalized, and reduce dataset
shuffle(vec(data))
data /= 255
data = data[1:50000,:]
# set architecture parameters
hidden_size = 500
latent_size = 25
n_components = 2
vae_params = init_params(size(data,2), hidden_size, latent_size, n_components)
# set hyperparams
adamParams = Dict("lr"=>0.0001, "m"=>init_params(size(data,2), hidden_size, latent_size, n_components, 0.),
"v"=>init_params(size(data,2), hidden_size, latent_size, n_components, 0.), "t"=>0)
hyperParams = Dict("adamParams"=>adamParams, "prior"=>Dict("weights"=>transpose([.5, .5]), "mu"=>[-2., 2.], "sigma"=>[1., 1.]), "nEpochs"=>50, "batchSize"=>100)
final_vae_params = trainVAE(data, vae_params, hyperParams)
end
main()