-
Notifications
You must be signed in to change notification settings - Fork 195
/
layers.py
853 lines (719 loc) · 43.1 KB
/
layers.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
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
### Custom Lasagne Layers for Introspective Adversarial Networks
# A Brock, 2016
#
# Layers that are not my own creation should be appropriately attributed here
# MADE wrapped from the implementation by M. Germain et al: https://github.com/mgermain/MADE
# Gaussian Sample layer from Tencia Lee's Recipe: https://github.com/Lasagne/Recipes/blob/master/examples/variational_autoencoder/variational_autoencoder.py
# Minibatch Discrimination layer from OpenAI's Improved GAN Techniques: https://github.com/openai/improved-gan
# Deconv Layer adapted from Radford's DCGAN: https://github.com/Newmu/dcgan_code
from __future__ import division
import numpy as np
import theano
import theano.tensor as T
import lasagne
import lasagne.layers
from lasagne.layers import SliceLayer as SL
from lasagne.layers import batch_norm as BN
from lasagne.layers import ElemwiseSumLayer as ESL
from lasagne.layers import NonlinearityLayer as NL
from lasagne.layers import DenseLayer as DL
from lasagne.init import Normal as initmethod
from lasagne.nonlinearities import elu
from lasagne.nonlinearities import rectify as relu
from lasagne.nonlinearities import LeakyRectify as lrelu
from lasagne.layers import TransposedConv2DLayer as TC2D
from lasagne.layers import ConcatLayer as CL
from math import sqrt
from theano.sandbox.rng_mrg import MRG_RandomStreams as RandomStreams
from mask_generator import MaskGenerator
# BatchReNorm Layer using cuDNN's BatchNorm
# This layer implements BatchReNorm (https://arxiv.org/abs/1702.03275),
# which modifies BatchNorm to include running-average statistics in addition to
# per-batch statistics in a well-principled manner. The RMAX and DMAX parameters
# are clip parameters which should be fscalars that you'll need to manage in the training
# loop, as they follow an annealing schedule (an example of which is given in the paper).
# I've been adjusting this schedule based on the total number of iterations relative
# to the number given in the paper, so for a ~50,000 iteration training run, I anneal
# RMAX between 1k and 5k iterations rather than 5k and 25k.
# NOTE: Ideally you should not have to manage RMAX and DMAX separately, so
# if someone wants to write a default_update similar to the one used for
# running_average and running_inv_std, that would be excellent.
class BatchReNormDNNLayer(lasagne.layers.BatchNormLayer):
def __init__(self, incoming, RMAX,DMAX,axes='auto', epsilon=1e-4, alpha=0.1,
beta=lasagne.init.Constant(0), gamma=lasagne.init.Constant(1),
mean=lasagne.init.Constant(0), inv_std=lasagne.init.Constant(1), **kwargs):
super(BatchReNormDNNLayer, self).__init__(
incoming, axes, epsilon, alpha, beta, gamma, mean, inv_std,
**kwargs)
all_but_second_axis = (0,) + tuple(range(2, len(self.input_shape)))
self.RMAX,self.DMAX = RMAX,DMAX
if self.axes not in ((0,), all_but_second_axis):
raise ValueError("BatchNormDNNLayer only supports normalization "
"across the first axis, or across all but the "
"second axis, got axes=%r" % (axes,))
def get_output_for(self, input, deterministic=False,
batch_norm_use_averages=None,
batch_norm_update_averages=None, **kwargs):
# Decide whether to use the stored averages or mini-batch statistics
if batch_norm_use_averages is None:
batch_norm_use_averages = deterministic
use_averages = batch_norm_use_averages
# Decide whether to update the stored averages
if batch_norm_update_averages is None:
batch_norm_update_averages = not deterministic
update_averages = batch_norm_update_averages
# prepare dimshuffle pattern inserting broadcastable axes as needed
param_axes = iter(range(input.ndim - len(self.axes)))
pattern = ['x' if input_axis in self.axes
else next(param_axes)
for input_axis in range(input.ndim)]
# and prepare the converse pattern removing those broadcastable axes
unpattern = [d for d in range(input.ndim) if d not in self.axes]
# call cuDNN if needed, obtaining normalized outputs and statistics
if not use_averages or update_averages:
# cuDNN requires beta/gamma tensors; create them if needed
shape = tuple(s for (d, s) in enumerate(input.shape)
if d not in self.axes)
gamma = self.gamma or theano.tensor.ones(shape)
beta = self.beta or theano.tensor.zeros(shape)
mode = 'per-activation' if self.axes == (0,) else 'spatial'
(normalized,
input_mean,
input_inv_std) = theano.sandbox.cuda.dnn.dnn_batch_normalization_train(
input, gamma.dimshuffle(pattern), beta.dimshuffle(pattern),
mode, self.epsilon)
# normalize with stored averages, if needed
if use_averages:
mean = self.mean.dimshuffle(pattern)
inv_std = self.inv_std.dimshuffle(pattern)
gamma = 1 if self.gamma is None else self.gamma.dimshuffle(pattern)
beta = 0 if self.beta is None else self.beta.dimshuffle(pattern)
normalized = (input - mean) * (gamma * inv_std) + beta
# update stored averages, if needed
if update_averages:
# Trick: To update the stored statistics, we create memory-aliased
# clones of the stored statistics:
running_mean = theano.clone(self.mean, share_inputs=False)
running_inv_std = theano.clone(self.inv_std, share_inputs=False)
# set a default update for them:
running_mean.default_update = ((1 - self.alpha) * running_mean +
self.alpha * input_mean.dimshuffle(unpattern))
running_inv_std.default_update = ((1 - self.alpha) *
running_inv_std +
self.alpha * input_inv_std.dimshuffle(unpattern))
# and make sure they end up in the graph without participating in
# the computation (this way their default_update will be collected
# and applied, but the computation will be optimized away):
# dummy = running_mean + running_inv_std).dimshuffle(pattern)
r = T.clip(running_inv_std.dimshuffle(pattern)/input_inv_std,1/self.RMAX,self.RMAX)
d = T.clip( (input_mean-running_mean.dimshuffle(pattern))*running_inv_std.dimshuffle(pattern),-self.DMAX,self.DMAX)
normalized = normalized * r + d
return normalized
# More Efficient MDCL layer
# When seeking to construct an MDC block, drop this into a Conv2D layer instead; it's faster.
# You can also easily re-parameterize this to a full-rank MDC block by dropping in the
# line that creates baseW into the for loop such that a new W is sampled each time.
def mdclW(num_filters,num_channels,filter_size,winit,name,scales):
# Coefficient Initializer
sinit = lasagne.init.Constant(1.0/(1+len(scales)))
# Total filter size
size = filter_size + (filter_size-1)*(scales[-1]-1)
# Multiscale Dilated Filter
W = T.zeros((num_filters,num_channels,size,size))
# Undilated Base Filter
baseW = theano.shared(lasagne.utils.floatX(winit.sample((num_filters,num_channels,filter_size,filter_size))),name=name+'.W')
for scale in enumerate(scales[::-1]): # enumerate backwards so that we place the main filter on top
W = T.set_subtensor(W[:,:,scales[-1]-scale:size-scales[-1]+scale:scale,scales[-1]-scale:size-scales[-1]+scale:scale],
baseW*theano.shared(lasagne.utils.floatX(sinit.sample(num_filters)), name+'.coeff_'+str(scale)).dimshuffle(0,'x','x','x'))
return W
# Subpixel Upsample Layer from (https://arxiv.org/abs/1609.05158)
# This layer uses a set of r^2 set_subtensor calls to reorganize the tensor in a subpixel-layer upscaling style
# as done in the ESPCN Magic ony paper for super-resolution.
# r is the upscale factor.
# c is the number of output channels.
class SubpixelLayer(lasagne.layers.Layer):
def __init__(self, incoming,r,c, **kwargs):
super(SubpixelLayer, self).__init__(incoming, **kwargs)
self.r=r # Upscale factor
self.c=c # number of output channels
def get_output_shape_for(self, input_shape):
return (input_shape[0],self.c,self.r*input_shape[2],self.r*input_shape[3])
def get_output_for(self, input, deterministic=False, **kwargs):
out = T.zeros((input.shape[0],self.output_shape[1],self.output_shape[2],self.output_shape[3]))
for x in xrange(self.r): # loop across all feature maps belonging to this channel
for y in xrange(self.r):
out=T.set_subtensor(out[:,:,x::self.r,y::self.r],input[:,self.r*x+y::self.r*self.r,:,:])
return out
# Subpixel Upsample Layer using reshapes as in https://github.com/Tetrachrome/subpixel. This implementation appears to be 10x slower than
# the set_subtensor implementation, presumably because of the extra reshapes after the splits.
class SubpixelLayer2(lasagne.layers.Layer):
def __init__(self, incoming,r,c, **kwargs):
super(SubpixelLayer2, self).__init__(incoming, **kwargs)
self.r=r
self.c=c
def get_output_shape_for(self, input_shape):
return (input_shape[0],self.c,self.r*input_shape[2],self.r*input_shape[3])
def get_output_for(self, input, deterministic=False, **kwargs):
def _phase_shift(input,r):
bsize,c,a,b = input.shape[0],1,self.output_shape[2]//r,self.output_shape[3]//r
X = T.reshape(input, (bsize,r,r,a,b))
X = T.transpose(X, (0, 3,4,1,2)) # bsize, a, b, r2,r1
X = T.split(x=X,splits_size=[1]*a,n_splits=a,axis=1) # a, [bsize, b, r, r]
X = [T.reshape(x,(bsize,b,r,r))for x in X]
X = T.concatenate(X,axis=2) # bsize, b, a*r, r
X = T.split(x=X,splits_size =[1]*b,n_splits=b,axis=1) # b, [bsize, a*r, r]
X = [T.reshape(x,(bsize,a*r,r))for x in X]
X = T.concatenate(X,axis=2) # bsize, a*r, b*r
return X.dimshuffle(0,'x',1,2)
Xc = T.split(x=input,splits_size =[input.shape[1]//self.c]*self.c,n_splits=self.c,axis=1)
return T.concatenate([_phase_shift(xc,self.r) for xc in Xc],axis=1)
# Multiscale Dilated Convolution Block
# This function (not a layer in and of itself, though you could make it one) returns a set of concatenated conv2d and dilatedconv2d layers.
# Each layer uses the same basic filter W, operating at a different dilation factor (or taken as the mean of W for the 1x1 conv).
# The channel-wise output of each layer is weighted by a set of coefficients, which are initialized to 1 / the total number of dilation scales,
# meaning that were starting by taking an elementwise mean. These should be learnable parameters.
# NOTES: - I'm considering changing the variable names to be more descriptive, and look less like ridiculous academic code. It's on the to-do list.
# - I keep the bias and nonlinearity out of the default definition for this layer, as I expect it to be batchnormed and nonlinearized in the model config.
def MDCL(incoming,num_filters,scales,name,dnn=True):
if dnn:
from lasagne.layers.dnn import Conv2DDNNLayer as C2D
# W initialization method--this should also work as Orthogonal('relu'), but I have yet to validate that as thoroughly.
winit = initmethod(0.02)
# Initialization method for the coefficients
sinit = lasagne.init.Constant(1.0/(1+len(scales)))
# Number of incoming channels
ni =lasagne.layers.get_output_shape(incoming)[1]
# Weight parameter--the primary parameter for this block
W = theano.shared(lasagne.utils.floatX(winit.sample((num_filters,lasagne.layers.get_output_shape(incoming)[1],3,3))),name=name+'W')
# Primary Convolution Layer--No Dilation
n = C2D(incoming = incoming,
num_filters = num_filters,
filter_size = [3,3],
stride = [1,1],
pad = (1,1),
W = W*theano.shared(lasagne.utils.floatX(sinit.sample(num_filters)), name+'_coeff_base').dimshuffle(0,'x','x','x'), # Note the broadcasting dimshuffle for the num_filter scalars.
b = None,
nonlinearity = None,
name = name+'base'
)
# List of remaining layers. This should probably just all be concatenated into a single list rather than being a separate deal.
nd = []
for i,scale in enumerate(scales):
# I don't think 0 dilation is technically defined (or if it is it's just the regular filter) but I use it here as a convenient keyword to grab the 1x1 mean conv.
if scale==0:
nd.append(C2D(incoming = incoming,
num_filters = num_filters,
filter_size = [1,1],
stride = [1,1],
pad = (0,0),
W = T.mean(W,axis=[2,3]).dimshuffle(0,1,'x','x')*theano.shared(lasagne.utils.floatX(sinit.sample(num_filters)), name+'_coeff_1x1').dimshuffle(0,'x','x','x'),
b = None,
nonlinearity = None,
name = name+str(scale)))
# Note the dimshuffles in this layer--these are critical as the current DilatedConv2D implementation uses a backward pass.
else:
nd.append(lasagne.layers.DilatedConv2DLayer(incoming = lasagne.layers.PadLayer(incoming = incoming, width=(scale,scale)),
num_filters = num_filters,
filter_size = [3,3],
dilation=(scale,scale),
W = W.dimshuffle(1,0,2,3)*theano.shared(lasagne.utils.floatX(sinit.sample(num_filters)), name+'_coeff_'+str(scale)).dimshuffle('x',0,'x','x'),
b = None,
nonlinearity = None,
name = name+str(scale)))
return ESL(nd+[n])
# MDC-based Upsample Layer.
# This is a prototype I don't make use of extensively. It's operational but it doesn't seem to improve results yet.
def USL(incoming,num_filters,scales,name,dnn=True):
if dnn:
from lasagne.layers.dnn import Conv2DDNNLayer as C2D
# W initialization method--this should also work as Orthogonal('relu'), but I have yet to validate that as thoroughly.
winit = initmethod(0.02)
# Initialization method for the coefficients
sinit = lasagne.init.Constant(1.0/(1+len(scales)))
# Number of incoming channels
ni =lasagne.layers.get_output_shape(incoming)[1]
# Weight parameter--the primary parameter for this block
W = theano.shared(lasagne.utils.floatX(winit.sample((num_filters,lasagne.layers.get_output_shape(incoming)[1],3,3))),name=name+'W')
# Primary Convolution Layer--No Dilation
n = C2D(incoming = Upscale2DLayer(incoming,2),
num_filters = num_filters,
filter_size = [3,3],
stride = [1,1],
pad = (1,1),
W = W*theano.shared(lasagne.utils.floatX(sinit.sample(num_filters)), name+'_coeff_base').dimshuffle(0,'x','x','x'),
b = None,
nonlinearity = None,
name = name+'base'
)
# Remaining layers
nd = []
for i,scale in enumerate(scales):
if scale==0:
nd.append(C2D(incoming = Upscale2DLayer(incoming,2),
num_filters = num_filters,
filter_size = [1,1],
stride = [1,1],
pad = (0,0),
W = T.mean(W,axis=[2,3]).dimshuffle(0,1,'x','x')*theano.shared(lasagne.utils.floatX(sinit.sample(num_filters)), name+'_coeff_1x1').dimshuffle(0,'x','x','x'),
b = None,
nonlinearity = None,
name = name+'1x1'
))
else:
nd.append(lasagne.layers.DilatedConv2DLayer(incoming = lasagne.layers.PadLayer(incoming = Upscale2DLayer(incoming,2), width=(scale,scale)),
num_filters = num_filters,
filter_size = [3,3],
dilation=(scale,scale),
W = W.dimshuffle(1,0,2,3)*theano.shared(lasagne.utils.floatX(sinit.sample(num_filters)), name+'_coeff_'+str(scale)).dimshuffle('x',0,'x','x'),
b = None,
nonlinearity = None,
name = name+str(scale)))
# A single deconv layer is also concatenated here. Like I said, it's a prototype!
nd.append(DeconvLayer(incoming = incoming,
num_filters = num_filters,
filter_size = [3,3],
stride = [2,2],
crop = (1,1),
W = W.dimshuffle(1,0,2,3)*theano.shared(lasagne.utils.floatX(sinit.sample(num_filters)), name+'_coeff_deconv').dimshuffle('x',0,'x','x'),
b = None,
nonlinearity = None,
name = name+'deconv'
))
return ESL(nd+[n])
#MDC-based Downsample Layer.
# This is a prototype I don't make use of extensively. It's operational and it seems like it works alright, but it's restrictively expensive
# and I am not PARALLELICUS, god of GPUs, so I don't have the memory to spare for it.
# Note that this layer does not currently support having a 0 scale like the others do, and just has a 1x1-stride2 conv by default.
def DSL(incoming,num_filters,scales,name,dnn=True):
if dnn:
from lasagne.layers.dnn import Conv2DDNNLayer as C2D
# W initialization method--this should also work as Orthogonal('relu'), but I have yet to validate that as thoroughly.
winit = initmethod(0.02)
# Initialization method for the coefficients
sinit = lasagne.init.Constant(1.0/(1+len(scales)))
# Number of incoming channels
ni =lasagne.layers.get_output_shape(incoming)[1]
# Weight parameter--the primary parameter for this block
W = theano.shared(lasagne.utils.floatX(winit.sample((num_filters,lasagne.layers.get_output_shape(incoming)[1],3,3))),name=name+'W')
# Main layer--3x3 conv with stride 2
n = C2D(incoming = incoming,
num_filters = num_filters,
filter_size = [3,3],
stride = [2,2],
pad = (1,1),
W = W*theano.shared(lasagne.utils.floatX(sinit.sample(num_filters)), name+'_coeff_base').dimshuffle(0,'x','x','x'),
b = None,
nonlinearity = None,
name = name+'base'
)
nd = []
for i,scale in enumerate(scales):
p = P2D(incoming = incoming,
pool_size = scale,
stride = 2,
pad = (1,1) if i else (0,0),
mode = 'average_exc_pad',
)
nd.append(C2D(incoming = p,
num_filters = num_filters,
filter_size = [3,3],
stride = (1,1),
pad = (1,1),
W = W*theano.shared(lasagne.utils.floatX(sinit.sample(num_filters)), name+'_coeff_'+str(scale)).dimshuffle(0,'x','x','x'),#.dimshuffle('x',0),
b = None,
nonlinearity = None,
name = name+str(scale)))
nd.append(C2D(incoming = incoming,
num_filters = num_filters,
filter_size = [1,1],
stride = [2,2],
pad = (0,0),
W = T.mean(W,axis=[2,3]).dimshuffle(0,1,'x','x')*theano.shared(lasagne.utils.floatX(sinit.sample(num_filters)), name+'_coeff_1x1').dimshuffle(0,'x','x','x'),
b = None,
nonlinearity = None,
name = name+'1x1'
))
return ESL(nd+[n])
# Beta Distribution Layer
# This layer takes in a batch_size batch, 2-channel, NxN dimension layer and returns the output of the first channel
# divided by the sum of both channels, which is equivalent to finding the expected value for a beta distribution.
# Note that this version of the layer scales to {-1,1} for compatibility with tanh.
class beta_layer(lasagne.layers.MergeLayer):
def __init__(self, alpha,beta, **kwargs):
super(beta_layer, self).__init__([alpha,beta], **kwargs)
def get_output_shape_for(self, input_shape):
print(input_shape)
return input_shape[0]
def get_output_for(self, inputs, deterministic=False, **kwargs):
alpha,beta = inputs
# return 2*T.true_div(alpha,T.add(alpha,beta)+1e-8)-1
return 2*(alpha/(alpha+beta+1e-8))-1
# Convenience Function to produce a residual pre-activation MDCL block
def MDBLOCK(incoming,num_filters,scales,name,nonlinearity):
return NL(BN(ESL([incoming,
MDCL(NL(BN(MDCL(NL(BN(incoming,name=name+'bnorm0'),nonlinearity),num_filters,scales,name),name=name+'bnorm1'),nonlinearity),
num_filters,
scales,
name+'2')]),name=name+'bnorm2'),nonlinearity)
# Gaussian Sample Layer for VAE from Tencia Lee
class GaussianSampleLayer(lasagne.layers.MergeLayer):
def __init__(self, mu, logsigma, rng=None, **kwargs):
self.rng = rng if rng else RandomStreams(lasagne.random.get_rng().randint(1,2147462579))
super(GaussianSampleLayer, self).__init__([mu, logsigma], **kwargs)
def get_output_shape_for(self, input_shapes):
return input_shapes[0]
def get_output_for(self, inputs, deterministic=False, **kwargs):
mu, logsigma = inputs
shape=(self.input_shapes[0][0] or inputs[0].shape[0],
self.input_shapes[0][1] or inputs[0].shape[1])
if deterministic:
return mu
return mu + T.exp(logsigma) * self.rng.normal(shape)
# DeconvLayer adapted from Radford's DCGAN Implementation
class DeconvLayer(lasagne.layers.conv.BaseConvLayer):
def __init__(self, incoming, num_filters, filter_size, stride=(1, 1),
crop=0, untie_biases=False,
W=initmethod(), b=lasagne.init.Constant(0.),
nonlinearity=lasagne.nonlinearities.rectify, flip_filters=False,
**kwargs):
super(DeconvLayer, self).__init__(
incoming, num_filters, filter_size, stride, crop, untie_biases,
W, b, nonlinearity, flip_filters, n=2, **kwargs)
# rename self.crop to self.pad
self.crop = self.pad
del self.pad
def get_W_shape(self):
num_input_channels = self.input_shape[1]
# first two sizes are swapped compared to a forward convolution
return (num_input_channels, self.num_filters) + self.filter_size
def get_output_shape_for(self, input_shape):
# when called from the constructor, self.crop is still called self.pad:
crop = getattr(self, 'crop', getattr(self, 'pad', None))
crop = crop if isinstance(crop, tuple) else (crop,) * self.n
batchsize = input_shape[0]
return(batchsize,self.num_filters)+(input_shape[2]*2,input_shape[3]*2)
# return ((batchsize, self.num_filters) +
# tuple(conv_input_length(input, filter, stride, p)
# for input, filter, stride, p
# in zip(input_shape[2:], self.filter_size,
# self.stride, crop)))
def convolve(self, input, **kwargs):
# Messy to have these imports here, but seems to allow for switching DNN off.
from theano.sandbox.cuda.basic_ops import (as_cuda_ndarray_variable,
host_from_gpu,
gpu_contiguous, HostFromGpu,
gpu_alloc_empty)
from theano.sandbox.cuda.dnn import GpuDnnConvDesc, GpuDnnConv, GpuDnnConvGradI, dnn_conv, dnn_pool
# Straight outta Radford
img = gpu_contiguous(input)
kerns = gpu_contiguous(self.W)
desc = GpuDnnConvDesc(border_mode=self.crop, subsample=self.stride,
conv_mode='conv')(gpu_alloc_empty(img.shape[0], kerns.shape[1], img.shape[2]*self.stride[0], img.shape[3]*self.stride[1]).shape, kerns.shape)
out = gpu_alloc_empty(img.shape[0], kerns.shape[1], img.shape[2]*self.stride[0], img.shape[3]*self.stride[1])
conved = GpuDnnConvGradI()(kerns, img, out, desc)
return conved
# Minibatch discrimination layer from OpenAI's improved GAN techniques
class MinibatchLayer(lasagne.layers.Layer):
def __init__(self, incoming, num_kernels, dim_per_kernel=5, theta=lasagne.init.Normal(0.05),
log_weight_scale=lasagne.init.Constant(0.), b=lasagne.init.Constant(-1.), **kwargs):
super(MinibatchLayer, self).__init__(incoming, **kwargs)
self.num_kernels = num_kernels
num_inputs = int(np.prod(self.input_shape[1:]))
self.theta = self.add_param(theta, (num_inputs, num_kernels, dim_per_kernel), name="theta")
self.log_weight_scale = self.add_param(log_weight_scale, (num_kernels, dim_per_kernel), name="log_weight_scale")
self.W = self.theta * (T.exp(self.log_weight_scale)/T.sqrt(T.sum(T.square(self.theta),axis=0))).dimshuffle('x',0,1)
self.b = self.add_param(b, (num_kernels,), name="b")
def get_output_shape_for(self, input_shape):
return (input_shape[0], np.prod(input_shape[1:])+self.num_kernels)
def get_output_for(self, input, init=False, **kwargs):
if input.ndim > 2:
# if the input has more than two dimensions, flatten it into a
# batch of feature vectors.
input = input.flatten(2)
activation = T.tensordot(input, self.W, [[1], [0]])
abs_dif = (T.sum(abs(activation.dimshuffle(0,1,2,'x') - activation.dimshuffle('x',1,2,0)),axis=2)
+ 1e6 * T.eye(input.shape[0]).dimshuffle(0,'x',1))
if init:
mean_min_abs_dif = 0.5 * T.mean(T.min(abs_dif, axis=2),axis=0)
abs_dif /= mean_min_abs_dif.dimshuffle('x',0,'x')
self.init_updates = [(self.log_weight_scale, self.log_weight_scale-T.log(mean_min_abs_dif).dimshuffle(0,'x'))]
f = T.sum(T.exp(-abs_dif),axis=2)
if init:
mf = T.mean(f,axis=0)
f -= mf.dimshuffle('x',0)
self.init_updates.append((self.b, -mf))
else:
f += self.b.dimshuffle('x',0)
return T.concatenate([input, f], axis=1)
# Convenience function to define an inception-style block
def InceptionLayer(incoming,param_dict,block_name):
branch = [0]*len(param_dict)
# Loop across branches
for i,dict in enumerate(param_dict):
for j,style in enumerate(dict['style']): # Loop up branch
branch[i] = C2D(
incoming = branch[i] if j else incoming,
num_filters = dict['num_filters'][j],
filter_size = dict['filter_size'][j],
pad = dict['pad'][j] if 'pad' in dict else None,
stride = dict['stride'][j],
W = initmethod('relu'),
nonlinearity = dict['nonlinearity'][j],
name = block_name+'_'+str(i)+'_'+str(j)) if style=='convolutional'\
else NL(lasagne.layers.dnn.Pool2DDNNLayer(
incoming=incoming if j == 0 else branch[i],
pool_size = dict['filter_size'][j],
mode = dict['mode'][j],
stride = dict['stride'][j],
pad = dict['pad'][j],
name = block_name+'_'+str(i)+'_'+str(j)),
nonlinearity = dict['nonlinearity'][j]) if style=='pool'\
else lasagne.layers.DilatedConv2DLayer(
incoming = lasagne.layers.PadLayer(incoming = incoming if j==0 else branch[i],width = dict['pad'][j]) if 'pad' in dict else incoming if j==0 else branch[i],
num_filters = dict['num_filters'][j],
filter_size = dict['filter_size'][j],
dilation = dict['dilation'][j],
# pad = dict['pad'][j] if 'pad' in dict else None,
W = initmethod('relu'),
nonlinearity = dict['nonlinearity'][j],
name = block_name+'_'+str(i)+'_'+str(j)) if style== 'dilation'\
else DL(
incoming = incoming if j==0 else branch[i],
num_units = dict['num_filters'][j],
W = initmethod('relu'),
b = None,
nonlinearity = dict['nonlinearity'][j],
name = block_name+'_'+str(i)+'_'+str(j))
# Apply Batchnorm
branch[i] = BN(branch[i],name = block_name+'_bnorm_'+str(i)+'_'+str(j)) if dict['bnorm'][j] else branch[i]
# Concatenate Sublayers
return CL(incomings=branch,name=block_name)
# Convenience function to define an inception-style block with upscaling
def InceptionUpscaleLayer(incoming,param_dict,block_name):
branch = [0]*len(param_dict)
# Loop across branches
for i,dict in enumerate(param_dict):
for j,style in enumerate(dict['style']): # Loop up branch
branch[i] = TC2D(
incoming = branch[i] if j else incoming,
num_filters = dict['num_filters'][j],
filter_size = dict['filter_size'][j],
crop = dict['pad'][j] if 'pad' in dict else None,
stride = dict['stride'][j],
W = initmethod('relu'),
nonlinearity = dict['nonlinearity'][j],
name = block_name+'_'+str(i)+'_'+str(j)) if style=='convolutional'\
else NL(
incoming = lasagne.layers.dnn.Pool2DDNNLayer(
incoming = lasagne.layers.Upscale2DLayer(
incoming=incoming if j == 0 else branch[i],
scale_factor = dict['stride'][j]),
pool_size = dict['filter_size'][j],
stride = [1,1],
mode = dict['mode'][j],
pad = dict['pad'][j],
name = block_name+'_'+str(i)+'_'+str(j)),
nonlinearity = dict['nonlinearity'][j])
# Apply Batchnorm
branch[i] = BN(branch[i],name = block_name+'_bnorm_'+str(i)+'_'+str(j)) if dict['bnorm'][j] else branch[i]
# Concatenate Sublayers
return CL(incomings=branch,name=block_name)
# Convenience function to efficiently generate param dictionaries for use with InceptioNlayer
def pd(num_layers=2,num_filters=32,filter_size=(3,3),pad=1,stride = (1,1),nonlinearity=elu,style='convolutional',bnorm=1,**kwargs):
input_args = locals()
input_args.pop('num_layers')
return {key:entry if type(entry) is list else [entry]*num_layers for key,entry in input_args.iteritems()}
# Possible Conv2DDNN convenience function. Remember to delete the C2D import at the top if you use this
# def C2D(incoming = None, num_filters = 32, filter_size= [3,3],pad = 'same',stride = [1,1], W = initmethod('relu'),nonlinearity = elu,name = None):
# return lasagne.layers.dnn.Conv2DDNNLayer(incoming,num_filters,filter_size,stride,pad,False,W,None,nonlinearity,False)
# Shape-Preserving Gaussian Sample layer for latent vectors with spatial dimensions.
# This is a holdover from an "old" (i.e. I abandoned it last month) idea.
class GSL(lasagne.layers.MergeLayer):
def __init__(self, mu, logsigma, rng=None, **kwargs):
self.rng = rng if rng else RandomStreams(lasagne.random.get_rng().randint(1,2147462579))
super(GSL, self).__init__([mu, logsigma], **kwargs)
def get_output_shape_for(self, input_shape):
print(input_shape)
return input_shape[0]
def get_output_for(self, inputs, deterministic=False, **kwargs):
mu, logsigma = inputs
if deterministic:
return mu
return mu + T.exp(logsigma) * self.rng.normal(logsigma.shape)
# Convenience function to return list of sampled latent layers
def GL(mu,ls):
return([GSL(z_mu,z_ls) for z_mu,z_ls in zip(mu,ls)])
# Convenience function to return a residual layer. It's not really that much more convenient than ESL'ing,
# but I like being able to see when I'm using Residual connections as opposed to Elemwise-sums
def ResLayer(incoming, IB,nonlinearity):
return NL(ESL([IB,incoming]),nonlinearity)
# Inverse autoregressive flow layer
class IAFLayer(lasagne.layers.MergeLayer):
def __init__(self, z, mu, logsigma, **kwargs):
super(IAFLayer, self).__init__([z,mu, logsigma], **kwargs)
def get_output_shape_for(self, input_shapes):
return input_shapes[0]
def get_output_for(self, inputs, deterministic=False, **kwargs):
z,mu, logsigma = inputs
return (z - mu) / T.exp(logsigma)
# Masked layer for MADE, adopted from M.Germain
class MaskedLayer(lasagne.layers.DenseLayer):
def __init__(self, incoming, num_units, mask_generator,layerIdx,W=lasagne.init.GlorotUniform(),
b=lasagne.init.Constant(0.), nonlinearity=lasagne.nonlinearities.rectify, **kwargs):
super(MaskedLayer, self).__init__(incoming, num_units, W,b, nonlinearity,**kwargs)
self.mask_generator = mask_generator
num_inputs = int(np.prod(self.input_shape[1:]))
self.weights_mask = self.add_param(spec = np.ones((num_inputs, num_units),dtype=np.float32),
shape = (num_inputs, num_units),
name='weights_mask',
trainable=False,
regularizable=False)
self.layerIdx = layerIdx
self.shuffle_update = [(self.weights_mask, mask_generator.get_mask_layer_UPDATE(self.layerIdx))]
def get_output_for(self,input, **kwargs):
if input.ndim > 2:
input = input.flatten(2)
activation = T.dot(input, self.W*self.weights_mask)
if self.b is not None:
activation = activation + self.b.dimshuffle('x', 0)
return self.nonlinearity(activation)
# Stripped-Down Direct Input masked layer: Combine this with ESL and a masked layer to get a true DIML.
# Consider making this a simultaneous subclass of MaskedLayer and elemwise sum layer for cleanliness
# adopted from M.Germain
class DIML(lasagne.layers.DenseLayer):
def __init__(self, incoming, num_units, mask_generator,layerIdx,W=lasagne.init.GlorotUniform(),
b=lasagne.init.Constant(0.), nonlinearity=None,**kwargs):
super(DIML, self).__init__(incoming, num_units, W,b, nonlinearity,**kwargs)
self.mask_generator = mask_generator
self.layerIdx = layerIdx
num_inputs = int(np.prod(self.input_shape[1:]))
self.weights_mask = self.add_param(spec = np.ones((num_inputs, num_units),dtype=np.float32),
shape = (num_inputs, num_units),
name='weights_mask',
trainable=False,
regularizable=False)
self.shuffle_update = [(self.weights_mask, self.mask_generator.get_direct_input_mask_layer_UPDATE(self.layerIdx + 1))]
def get_output_for(self,input, **kwargs):
if input.ndim > 2:
input = input.flatten(2)
activation = T.dot(input, self.W*self.weights_mask)
if self.b is not None:
activation = activation + self.b.dimshuffle('x', 0)
return self.nonlinearity(activation)
# Conditioning Masked Layer
# Currently not used.
# class CML(MaskedLayer):
# def __init__(self, incoming, num_units, mask_generator,use_cond_mask=False,U=lasagne.init.GlorotUniform(),W=lasagne.init.GlorotUniform(),
# b=init.Constant(0.), nonlinearity=lasagne.nonlinearities.rectify, **kwargs):
# super(CML, self).__init__(incoming, num_units, mask_generator,W,
# b, nonlinearity,**kwargs)
# self.use_cond_mask=use_cond_mask
# if use_cond_mask:
# self.U = self.add_param(spec = U,
# shape = (num_inputs, num_units),
# name='U',
# trainable=True,
# regularizable=False)theano.shared(value=self.weights_initialization((self.n_in, self.n_out)), name=self.name+'U', borrow=True)
# self.add_param(self.U,name =
# def get_output_for(self,input,**kwargs):
# lin = self.lin_output = T.dot(input, self.W * self.weights_mask) + self.b
# if self.use_cond_mask:
# lin = lin+T.dot(T.ones_like(input), self.U * self.weights_mask)
# return lin if self._activation is None else self._activation(lin)
# Made layer, adopted from M.Germain
class MADE(lasagne.layers.Layer):
def __init__(self,z,hidden_sizes,name,nonlinearity=lasagne.nonlinearities.rectify,output_nonlinearity=None, **kwargs):
# self.rng = rng if rng else RandomStreams(lasagne.random.get_rng().randint(1234))
super(MADE, self).__init__(z, **kwargs)
# Incoming latents
self.z = z
# List defining hidden units in each layer
self.hidden_sizes = hidden_sizes
# Layer name for saving parameters.
self.name = name
# nonlinearity
self.nonlinearity = nonlinearity
# Output nonlinearity
self.output_nonlinearity = output_nonlinearity
# Control parameters from original MADE
mask_distribution=0
use_cond_mask = False
direct_input_connect = "Output"
direct_output_connect = False
self.shuffled_once = False
# Mask generator
self.mask_generator = MaskGenerator(lasagne.layers.get_output_shape(z)[1], hidden_sizes, mask_distribution)
# Build the MADE
# TODO: Consider making this more compact by directly writing to the layers list
self.input_layer = MaskedLayer(incoming = z,
num_units = hidden_sizes[0],
mask_generator = self.mask_generator,
layerIdx = 0,
W = lasagne.init.Orthogonal('relu'),
nonlinearity=self.nonlinearity,
name = self.name+'_input')
self.layers = [self.input_layer]
for i in range(1, len(hidden_sizes)):
self.layers += [MaskedLayer(incoming = self.layers[-1],
num_units = hidden_sizes[i],
mask_generator = self.mask_generator,
layerIdx = i,
W = lasagne.init.Orthogonal('relu'),
nonlinearity=self.nonlinearity,
name = self.name+'_layer_'+str(i))]
outputLayerIdx = len(self.layers)
# Output layer
self.layers += [MaskedLayer(incoming = self.layers[-1],
num_units = lasagne.layers.get_output_shape(z)[1],
mask_generator = self.mask_generator,
layerIdx = outputLayerIdx,
W = lasagne.init.Orthogonal('relu'),
nonlinearity = self.output_nonlinearity,
name = self.name+'_output_W'),
DIML(incoming = z,
num_units = lasagne.layers.get_output_shape(z)[1],
mask_generator = self.mask_generator,
layerIdx = outputLayerIdx,
W = lasagne.init.Orthogonal('relu'),
nonlinearity = self.output_nonlinearity,
name = self.name+'_output_D')]
masks_updates = [layer_mask_update for l in self.layers for layer_mask_update in l.shuffle_update]
self.update_masks = theano.function(name='update_masks',
inputs=[],
updates=masks_updates)
# Make the true output layer by ESL'ing the DIML and masked layer
self.final_layer= ESL([self.layers[-2],self.layers[-1]])
# self.output_layer = self.layers[-1]
# params = [p for p in l.get_params(trainable=True) for l in self.layers]
# print(params)
def get_output_for(self, input, deterministic=False, **kwargs):
return lasagne.layers.get_output(self.final_layer,{self.z:input})
def get_params(self, unwrap_shared=True, **tags):
params = []
for l in self.layers:
for p in l.get_params(**tags):
params.append(p)
return(params)
# params = [p for p in l.get_params(trainable=True) for l in self.layers]
# return params
# return [p for p in lay.get_params(unwrap_shared,**tags) for lay in self.layers]
# return lasagne.layers.get_all_params(self.final_layer,trainable=True)
def shuffle(self, shuffling_type):
if shuffling_type == "Once" and self.shuffled_once is False:
self.mask_generator.shuffle_ordering()
self.mask_generator.sample_connectivity()
self.update_masks()
self.shuffled_once = True
return
if shuffling_type in ["Ordering", "Full"]:
self.mask_generator.shuffle_ordering()
if shuffling_type in ["Connectivity", "Full"]:
self.mask_generator.sample_connectivity()
self.update_masks()
def reset(self, shuffling_type, last_shuffle=0):
self.mask_generator.reset()
# Always do a first shuffle so that the natural order does not gives us an edge
self.shuffle("Full")
# Set the mask to the requested shuffle
for i in range(last_shuffle):
self.shuffle(shuffling_type)