Minor syntactic change.
[pytorch] / attentiontoy1d.py
1 #!/usr/bin/env python
2
3 # Any copyright is dedicated to the Public Domain.
4 # https://creativecommons.org/publicdomain/zero/1.0/
5
6 # Written by Francois Fleuret <francois@fleuret.org>
7
8 import torch, math, sys, argparse
9
10 from torch import nn
11 from torch.nn import functional as F
12
13 import matplotlib.pyplot as plt
14
15 ######################################################################
16
17 parser = argparse.ArgumentParser(description='Toy attention model.')
18
19 parser.add_argument('--nb_epochs',
20                     type = int, default = 250)
21
22 parser.add_argument('--with_attention',
23                     help = 'Use the model with an attention layer',
24                     action='store_true', default=False)
25
26 parser.add_argument('--group_by_locations',
27                     help = 'Use the task where the grouping is location-based',
28                     action='store_true', default=False)
29
30 parser.add_argument('--positional_encoding',
31                     help = 'Provide a positional encoding',
32                     action='store_true', default=False)
33
34 parser.add_argument('--seed',
35                     type = int, default = 0,
36                     help = 'Random seed (default 0, < 0 is no seeding)')
37
38 args = parser.parse_args()
39
40 if args.seed >= 0:
41     torch.manual_seed(args.seed)
42
43 ######################################################################
44
45 label=''
46
47 if args.with_attention: label = 'wa_'
48
49 if args.group_by_locations: label += 'lg_'
50
51 if args.positional_encoding: label += 'pe_'
52
53 log_file = open(f'att1d_{label}train.log', 'w')
54
55 ######################################################################
56
57 def log_string(s):
58     if log_file is not None:
59         log_file.write(s + '\n')
60         log_file.flush()
61     print(s)
62     sys.stdout.flush()
63
64 ######################################################################
65
66 if torch.cuda.is_available():
67     device = torch.device('cuda')
68     torch.backends.cudnn.benchmark = True
69 else:
70     device = torch.device('cpu')
71
72 ######################################################################
73
74 seq_height_min, seq_height_max = 1.0, 25.0
75 seq_width_min, seq_width_max = 5.0, 11.0
76 seq_length = 100
77
78 def positions_to_sequences(tr = None, bx = None, noise_level = 0.3):
79     st = torch.arange(seq_length).float()
80     st = st[None, :, None]
81     tr = tr[:, None, :, :]
82     bx = bx[:, None, :, :]
83
84     xtr =            torch.relu(tr[..., 1] - torch.relu(torch.abs(st - tr[..., 0]) - 0.5) * 2 * tr[..., 1] / tr[..., 2])
85     xbx = torch.sign(torch.relu(bx[..., 1] - torch.abs((st - bx[..., 0]) * 2 * bx[..., 1] / bx[..., 2]))) * bx[..., 1]
86
87     x = torch.cat((xtr, xbx), 2)
88
89     # u = x.sign()
90     u = F.max_pool1d(x.sign().permute(0, 2, 1), kernel_size = 2, stride = 1).permute(0, 2, 1)
91
92     collisions = (u.sum(2) > 1).max(1).values
93     y = x.max(2).values
94
95     return y + torch.rand_like(y) * noise_level - noise_level / 2, collisions
96
97 ######################################################################
98
99 def generate_sequences(nb):
100
101     # Position / height / width
102
103     tr = torch.empty(nb, 2, 3)
104     tr[:, :, 0].uniform_(seq_width_max/2, seq_length - seq_width_max/2)
105     tr[:, :, 1].uniform_(seq_height_min, seq_height_max)
106     tr[:, :, 2].uniform_(seq_width_min, seq_width_max)
107
108     bx = torch.empty(nb, 2, 3)
109     bx[:, :, 0].uniform_(seq_width_max/2, seq_length - seq_width_max/2)
110     bx[:, :, 1].uniform_(seq_height_min, seq_height_max)
111     bx[:, :, 2].uniform_(seq_width_min, seq_width_max)
112
113     if args.group_by_locations:
114         a = torch.cat((tr, bx), 1)
115         v = a[:, :, 0].sort(1).values[:, 2:3]
116         mask_left = (a[:, :, 0] < v).float()
117         h_left = (a[:, :, 1] * mask_left).sum(1) / 2
118         h_right = (a[:, :, 1] * (1 - mask_left)).sum(1) / 2
119         valid = (h_left - h_right).abs() > 4
120     else:
121         valid = (torch.abs(tr[:, 0, 1] - tr[:, 1, 1]) > 4) & (torch.abs(tr[:, 0, 1] - tr[:, 1, 1]) > 4)
122
123     input, collisions = positions_to_sequences(tr, bx)
124
125     if args.group_by_locations:
126         a = torch.cat((tr, bx), 1)
127         v = a[:, :, 0].sort(1).values[:, 2:3]
128         mask_left = (a[:, :, 0] < v).float()
129         h_left = (a[:, :, 1] * mask_left).sum(1, keepdim = True) / 2
130         h_right = (a[:, :, 1] * (1 - mask_left)).sum(1, keepdim = True) / 2
131         a[:, :, 1] = mask_left * h_left + (1 - mask_left) * h_right
132         tr, bx = a.split(2, 1)
133     else:
134         tr[:, :, 1:2] = tr[:, :, 1:2].mean(1, keepdim = True)
135         bx[:, :, 1:2] = bx[:, :, 1:2].mean(1, keepdim = True)
136
137     targets, _ = positions_to_sequences(tr, bx)
138
139     valid = valid & ~collisions
140     tr = tr[valid]
141     bx = bx[valid]
142     input = input[valid][:, None, :]
143     targets = targets[valid][:, None, :]
144
145     if input.size(0) < nb:
146         input2, targets2, tr2, bx2 = generate_sequences(nb - input.size(0))
147         input = torch.cat((input, input2), 0)
148         targets = torch.cat((targets, targets2), 0)
149         tr = torch.cat((tr, tr2), 0)
150         bx = torch.cat((bx, bx2), 0)
151
152     return input, targets, tr, bx
153
154 ######################################################################
155
156 def save_sequence_images(filename, sequences, tr = None, bx = None):
157     fig = plt.figure()
158     ax = fig.add_subplot(1, 1, 1)
159
160     ax.set_xlim(0, seq_length)
161     ax.set_ylim(-1, seq_height_max + 4)
162
163     for u in sequences:
164         ax.plot(
165             torch.arange(u[0].size(0)) + 0.5, u[0], color = u[1], label = u[2]
166         )
167
168     ax.legend(frameon = False, loc = 'upper left')
169
170     delta = -1.
171     if tr is not None:
172         ax.scatter(test_tr[k, :, 0], torch.full((test_tr.size(1),), delta), color = 'black', marker = '^', clip_on=False)
173
174     if bx is not None:
175         ax.scatter(test_bx[k, :, 0], torch.full((test_bx.size(1),), delta), color = 'black', marker = 's', clip_on=False)
176
177     fig.savefig(filename, bbox_inches='tight')
178
179     plt.close('all')
180
181 ######################################################################
182
183 class AttentionLayer(nn.Module):
184     def __init__(self, in_channels, out_channels, key_channels):
185         super(AttentionLayer, self).__init__()
186         self.conv_Q = nn.Conv1d(in_channels, key_channels, kernel_size = 1, bias = False)
187         self.conv_K = nn.Conv1d(in_channels, key_channels, kernel_size = 1, bias = False)
188         self.conv_V = nn.Conv1d(in_channels, out_channels, kernel_size = 1, bias = False)
189
190     def forward(self, x):
191         Q = self.conv_Q(x)
192         K = self.conv_K(x)
193         V = self.conv_V(x)
194         A = Q.permute(0, 2, 1).matmul(K).softmax(2)
195         x = A.matmul(V.permute(0, 2, 1)).permute(0, 2, 1)
196         return x
197
198     def __repr__(self):
199         return self._get_name() + \
200             '(in_channels={}, out_channels={}, key_channels={})'.format(
201                 self.conv_Q.in_channels,
202                 self.conv_V.out_channels,
203                 self.conv_K.out_channels
204             )
205
206     def attention(self, x):
207         Q = self.conv_Q(x)
208         K = self.conv_K(x)
209         return Q.permute(0, 2, 1).matmul(K).softmax(2)
210
211 ######################################################################
212
213 train_input, train_targets, train_tr, train_bx = generate_sequences(25000)
214 test_input, test_targets, test_tr, test_bx = generate_sequences(1000)
215
216 ######################################################################
217
218 ks = 5
219 nc = 64
220
221 if args.positional_encoding:
222     c = math.ceil(math.log(seq_length) / math.log(2.0))
223     positional_input = (torch.arange(seq_length).unsqueeze(0) // 2**torch.arange(c).unsqueeze(1))%2
224     positional_input = positional_input.unsqueeze(0).float()
225 else:
226     positional_input = torch.zeros(1, 0, seq_length)
227
228 in_channels = 1 + positional_input.size(1)
229
230 if args.with_attention:
231
232     model = nn.Sequential(
233         nn.Conv1d(in_channels, nc, kernel_size = ks, padding = ks//2),
234         nn.ReLU(),
235         nn.Conv1d(nc, nc, kernel_size = ks, padding = ks//2),
236         nn.ReLU(),
237         AttentionLayer(nc, nc, nc),
238         nn.Conv1d(nc, nc, kernel_size = ks, padding = ks//2),
239         nn.ReLU(),
240         nn.Conv1d(nc,  1, kernel_size = ks, padding = ks//2)
241     )
242
243 else:
244
245     model = nn.Sequential(
246         nn.Conv1d(in_channels, nc, kernel_size = ks, padding = ks//2),
247         nn.ReLU(),
248         nn.Conv1d(nc, nc, kernel_size = ks, padding = ks//2),
249         nn.ReLU(),
250         nn.Conv1d(nc, nc, kernel_size = ks, padding = ks//2),
251         nn.ReLU(),
252         nn.Conv1d(nc, nc, kernel_size = ks, padding = ks//2),
253         nn.ReLU(),
254         nn.Conv1d(nc,  1, kernel_size = ks, padding = ks//2)
255     )
256
257 nb_parameters = sum(p.numel() for p in model.parameters())
258
259 with open(f'att1d_{label}model.log', 'w') as f:
260     f.write(str(model) + '\n\n')
261     f.write(f'nb_parameters {nb_parameters}\n')
262
263 ######################################################################
264
265 batch_size = 100
266
267 optimizer = torch.optim.Adam(model.parameters(), lr = 1e-3)
268 mse_loss = nn.MSELoss()
269
270 model.to(device)
271 mse_loss.to(device)
272 train_input, train_targets = train_input.to(device), train_targets.to(device)
273 test_input, test_targets = test_input.to(device), test_targets.to(device)
274 positional_input = positional_input.to(device)
275
276 mu, std = train_input.mean(), train_input.std()
277
278 for e in range(args.nb_epochs):
279     acc_loss = 0.0
280
281     for input, targets in zip(train_input.split(batch_size),
282                               train_targets.split(batch_size)):
283
284         input = torch.cat((input, positional_input.expand(input.size(0), -1, -1)), 1)
285
286         output = model((input - mu) / std)
287         loss = mse_loss(output, targets)
288
289         optimizer.zero_grad()
290         loss.backward()
291         optimizer.step()
292
293         acc_loss += loss.item()
294
295     log_string(f'{e+1} {acc_loss}')
296
297 ######################################################################
298
299 train_input = train_input.detach().to('cpu')
300 train_targets = train_targets.detach().to('cpu')
301
302 for k in range(15):
303     save_sequence_images(
304         f'att1d_{label}train_{k:03d}.pdf',
305         [
306             ( train_input[k, 0], 'blue', 'Input' ),
307             ( train_targets[k, 0], 'red', 'Target' ),
308         ],
309     )
310
311 ####################
312
313 test_input = torch.cat((test_input, positional_input.expand(test_input.size(0), -1, -1)), 1)
314 test_outputs = model((test_input - mu) / std).detach()
315
316 if args.with_attention:
317     k = next(k for k, l in enumerate(model) if isinstance(l, AttentionLayer))
318     x = model[0:k]((test_input - mu) / std)
319     test_A = model[k].attention(x)
320     test_A = test_A.detach().to('cpu')
321
322 test_input = test_input.detach().to('cpu')
323 test_outputs = test_outputs.detach().to('cpu')
324 test_targets = test_targets.detach().to('cpu')
325
326 for k in range(15):
327     save_sequence_images(
328         f'att1d_{label}test_Y_{k:03d}.pdf',
329         [
330             ( test_input[k, 0], 'blue', 'Input' ),
331             ( test_outputs[k, 0], 'orange', 'Output' ),
332         ]
333     )
334
335     save_sequence_images(
336         f'att1d_{label}test_Yp_{k:03d}.pdf',
337         [
338             ( test_input[k, 0], 'blue', 'Input' ),
339             ( test_outputs[k, 0], 'orange', 'Output' ),
340         ],
341         test_tr[k],
342         test_bx[k]
343     )
344
345     if args.with_attention:
346         fig = plt.figure()
347         ax = fig.add_subplot(1, 1, 1)
348         ax.set_xlim(0, seq_length)
349         ax.set_ylim(0, seq_length)
350
351         ax.imshow(test_A[k], cmap = 'binary', interpolation='nearest')
352         delta = 0.
353         ax.scatter(test_bx[k, :, 0], torch.full((test_bx.size(1),), delta), color = 'black', marker = 's', clip_on=False)
354         ax.scatter(torch.full((test_bx.size(1),), delta), test_bx[k, :, 0], color = 'black', marker = 's', clip_on=False)
355         ax.scatter(test_tr[k, :, 0], torch.full((test_tr.size(1),), delta), color = 'black', marker = '^', clip_on=False)
356         ax.scatter(torch.full((test_tr.size(1),), delta), test_tr[k, :, 0], color = 'black', marker = '^', clip_on=False)
357
358         fig.savefig(f'att1d_{label}test_A_{k:03d}.pdf', bbox_inches='tight')
359
360     plt.close('all')
361
362 ######################################################################