Update.
[pytorch.git] / autoencoder.py
1 #!/usr/bin/env python
2
3 import sys, argparse, os, time
4
5 import torch, torchvision
6
7 from torch import optim, nn
8 from torch.nn import functional as F
9
10 import torchvision
11
12 ######################################################################
13
14 if torch.cuda.is_available():
15     device = torch.device('cuda')
16 else:
17     device = torch.device('cpu')
18
19 ######################################################################
20
21 parser = argparse.ArgumentParser(description = 'Simple auto-encoder.')
22
23 parser.add_argument('--nb_epochs',
24                     type = int, default = 25)
25
26 parser.add_argument('--batch_size',
27                     type = int, default = 100)
28
29 parser.add_argument('--data_dir',
30                     type = str, default = './data/')
31
32 parser.add_argument('--log_filename',
33                     type = str, default = 'train.log')
34
35 parser.add_argument('--embedding_dim',
36                     type = int, default = 16)
37
38 parser.add_argument('--nb_channels',
39                     type = int, default = 32)
40
41 parser.add_argument('--force_train',
42                     type = bool, default = False)
43
44 args = parser.parse_args()
45
46 log_file = open(args.log_filename, 'w')
47
48 ######################################################################
49
50 def log_string(s, color = None):
51     t = time.strftime("%Y-%m-%d_%H:%M:%S - ", time.localtime())
52
53     if log_file is not None:
54         log_file.write(t + s + '\n')
55         log_file.flush()
56
57     print(t + s)
58     sys.stdout.flush()
59
60 ######################################################################
61
62 class AutoEncoder(nn.Module):
63     def __init__(self, nb_channels, embedding_dim):
64         super(AutoEncoder, self).__init__()
65
66         self.encoder = nn.Sequential(
67             nn.Conv2d(1, nb_channels, kernel_size = 5), # to 24x24
68             nn.ReLU(inplace = True),
69             nn.Conv2d(nb_channels, nb_channels, kernel_size = 5), # to 20x20
70             nn.ReLU(inplace = True),
71             nn.Conv2d(nb_channels, nb_channels, kernel_size = 4, stride = 2), # to 9x9
72             nn.ReLU(inplace = True),
73             nn.Conv2d(nb_channels, nb_channels, kernel_size = 3, stride = 2), # to 4x4
74             nn.ReLU(inplace = True),
75             nn.Conv2d(nb_channels, embedding_dim, kernel_size = 4)
76         )
77
78         self.decoder = nn.Sequential(
79             nn.ConvTranspose2d(embedding_dim, nb_channels, kernel_size = 4),
80             nn.ReLU(inplace = True),
81             nn.ConvTranspose2d(nb_channels, nb_channels, kernel_size = 3, stride = 2), # from 4x4
82             nn.ReLU(inplace = True),
83             nn.ConvTranspose2d(nb_channels, nb_channels, kernel_size = 4, stride = 2), # from 9x9
84             nn.ReLU(inplace = True),
85             nn.ConvTranspose2d(nb_channels, nb_channels, kernel_size = 5), # from 20x20
86             nn.ReLU(inplace = True),
87             nn.ConvTranspose2d(nb_channels, 1, kernel_size = 5), # from 24x24
88         )
89
90     def encode(self, x):
91         return self.encoder(x).view(x.size(0), -1)
92
93     def decode(self, z):
94         return self.decoder(z.view(z.size(0), -1, 1, 1))
95
96     def forward(self, x):
97         x = self.encoder(x)
98         # print(x.size())
99         x = self.decoder(x)
100         return x
101
102 ######################################################################
103
104 train_set = torchvision.datasets.MNIST(args.data_dir + '/mnist/',
105                                        train = True, download = True)
106 train_input = train_set.data.view(-1, 1, 28, 28).float()
107
108 test_set = torchvision.datasets.MNIST(args.data_dir + '/mnist/',
109                                       train = False, download = True)
110 test_input = test_set.data.view(-1, 1, 28, 28).float()
111
112 ######################################################################
113
114 train_input, test_input = train_input.to(device), test_input.to(device)
115
116 mu, std = train_input.mean(), train_input.std()
117 train_input.sub_(mu).div_(std)
118 test_input.sub_(mu).div_(std)
119
120 model = AutoEncoder(args.nb_channels, args.embedding_dim)
121 optimizer = optim.Adam(model.parameters(), lr = 1e-3)
122
123 model.to(device)
124
125 for epoch in range(args.nb_epochs):
126     acc_loss = 0
127     for input in train_input.split(args.batch_size):
128         input = input.to(device)
129         z = model.encode(input)
130         output = model.decode(z)
131         loss = 0.5 * (output - input).pow(2).sum() / input.size(0)
132
133         optimizer.zero_grad()
134         loss.backward()
135         optimizer.step()
136
137         acc_loss += loss.item()
138
139     log_string(f'acc_loss {epoch} {acc_loss}', 'blue')
140
141 ######################################################################
142
143 input = test_input[:256]
144 z = model.encode(input)
145 output = model.decode(z)
146
147 torchvision.utils.save_image(1 - input, 'ae-input.png', nrow = 16, pad_value = 0.8)
148 torchvision.utils.save_image(1 - output, 'ae-output.png', nrow = 16, pad_value = 0.8)
149
150 ######################################################################
151
152 input = train_input[:256]
153 z = model.encode(input)
154 mu, std = z.mean(0), z.std(0)
155 z = z.normal_() * std + mu
156 output = model.decode(z)
157 torchvision.utils.save_image(1 - output, 'ae-synth.png', nrow = 16, pad_value = 0.8)
158
159 ######################################################################