def make_madry_ngpu(nb_classes=10, input_shape=(None, 28, 28, 1), **kwargs):
"""
Create a multi-GPU model similar to Madry et al. (arXiv:1706.06083).
"""
layers = [Conv2DnGPU(32, (5, 5), (1, 1), "SAME"),
ReLU(),
MaxPool((2, 2), (2, 2), "SAME"),
Conv2DnGPU(64, (5, 5), (1, 1), "SAME"),
ReLU(),
MaxPool((2, 2), (2, 2), "SAME"),
Flatten(),
LinearnGPU(1024),
ReLU(),
LinearnGPU(nb_classes),
Softmax()]
model = MLPnGPU(layers, input_shape)
return model
def _conv(self, name, x, filter_size, in_filters, out_filters, strides):
"""Convolution."""
if self.init_layers:
conv = Conv2DnGPU(out_filters,
(filter_size, filter_size),
strides[1:3], 'SAME', w_name='DW')
conv.name = name
self.layers += [conv]
else:
conv = self.layers[self.layer_idx]
self.layer_idx += 1
conv.device_name = self.device_name
conv.set_training(self.training)
return conv.fprop(x)