-
Notifications
You must be signed in to change notification settings - Fork 0
/
convAE.py
67 lines (54 loc) · 2.56 KB
/
convAE.py
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
import mxnet as mx
from mxnet import gluon
import utils
# convolutional autoencoder
class ConvolutionalAutoencoder(gluon.nn.HybridBlock):
def __init__(self):
super(ConvolutionalAutoencoder, self).__init__()
with self.name_scope():
self.encoder = gluon.nn.HybridSequential(prefix="")
with self.encoder.name_scope():
self.encoder.add(gluon.nn.Conv2D(32, 5, padding=0, activation='relu'))
self.encoder.add(gluon.nn.MaxPool2D(2))
self.encoder.add(gluon.nn.Conv2D(32, 5, padding=0, activation='relu'))
self.encoder.add(gluon.nn.MaxPool2D(2))
self.encoder.add(gluon.nn.Dense(2000))
self.decoder = gluon.nn.HybridSequential(prefix="")
with self.decoder.name_scope():
self.decoder.add(gluon.nn.Dense(32*22*22, activation='relu'))
self.decoder.add(gluon.nn.HybridLambda(lambda F, x: F.UpSampling(x, scale=2, sample_type='nearest')))
self.decoder.add(gluon.nn.Conv2DTranspose(32, 5, activation='relu'))
self.decoder.add(gluon.nn.HybridLambda(lambda F, x: F.UpSampling(x, scale=2, sample_type='nearest')))
self.decoder.add(gluon.nn.Conv2DTranspose(1, kernel_size=5, activation='sigmoid'))
def hybrid_forward(self, F, x):
x = self.encoder(x)
x = self.decoder[0](x)
# need to reshape ouput feature vector from Dense(32*22*22), before it is upsampled
x = x.reshape((-1,32,22,22))
x = self.decoder[1:](x)
return x
# Train the autoencoder
def train(batch_size, ctx, num_epochs, path, lr=1e-4, wd=1e-5, params_file="autoencoder_ucsd_convae.params"):
# Dataloader for training dataset
dataloader = utils.create_dataset(path, batch_size, shuffle=True)
# Get model
model = ConvolutionalAutoencoder()
model.hybridize()
# Initialiize
model.collect_params().initialize(mx.init.Xavier('gaussian'), ctx=ctx)
# Loss
l2loss = gluon.loss.L2Loss()
optimizer = gluon.Trainer(model.collect_params(), 'adam', {'learning_rate': lr, 'wd': wd})
# Start training loop
for epoch in range(num_epochs):
for image in dataloader:
image = image.as_in_context(ctx)
with mx.autograd.record():
reconstructed = model(image)
loss = l2loss(reconstructed, image)
loss.backward()
optimizer.step(batch_size)
print('epoch [{}/{}], loss:{:.4f}'.format(epoch + 1, num_epochs, mx.nd.mean(loss).asscalar()))
# Save parameters
model.save_parameters(params_file)
return model, params_file