1from __future__ import absolute_import, print_function, division
2
3import theano
4from theano import Apply
5from theano.gof import ParamsType
6from theano.scalar import bool as bool_t
7from theano.tensor.basic import as_tensor_variable
8from theano.tensor.signal.pool import Pool, PoolingMode_t
9
10from .type import gpu_context_type
11from .basic_ops import (CGpuKernelBase, infer_context_name, gpuarray_helper_inc_dir,
12                        as_gpuarray_variable, gpu_contiguous)
13
14try:
15    import pygpu
16except ImportError as e:
17    # To make sure theano is importable
18    pass
19
20
21class GpuPool(CGpuKernelBase):
22    """
23    Implement the max and average pooling on the gpu.
24
25    """
26    __props__ = ('ignore_border', 'mode', 'ndim')
27    params_type = ParamsType(ignore_border=bool_t,
28                             mode=PoolingMode_t,
29                             context=gpu_context_type)
30
31    def __init__(self, ignore_border, mode='max', ndim=2):
32        self.ndim = ndim
33        self.ignore_border = ignore_border
34        if mode == 'average':
35            mode = 'average_inc_pad'
36        self.mode = mode
37        CGpuKernelBase.__init__(self, ['c_code/pool.c'],
38                                'APPLY_SPECIFIC(pool)')
39        assert PoolingMode_t.has_alias(self.mode)
40        assert self.ndim in [2, 3]
41
42    def get_params(self, node):
43        return self.params_type.get_params(self, context=node.inputs[0].type.context)
44
45    def c_headers(self):
46        return ['gpuarray_api.h', 'gpuarray_helper.h', 'numpy_compat.h']
47
48    def c_header_dirs(self):
49        return [gpuarray_helper_inc_dir(), pygpu.get_include()]
50
51    def make_node(self, inp, ws, stride=None, pad=None):
52        ctx_name = infer_context_name(inp)
53        inp = as_gpuarray_variable(inp, ctx_name)
54        nd = self.ndim
55        assert (inp.ndim == nd + 2)
56        if stride is None:
57            stride = ws
58        if pad is None:
59            pad = (0,) * nd
60        elif isinstance(pad, (tuple, list)):
61            if max(pad) != 0 and not self.ignore_border:
62                raise ValueError('Padding works only with ignore_border=True')
63            if isinstance(ws, (tuple, list)):
64                if any(pad[i] >= ws[i] for i in range(nd)):
65                    raise ValueError('Padding must be smaller than strides')
66
67        ws = as_tensor_variable(ws)
68        stride = as_tensor_variable(stride)
69        pad = as_tensor_variable(pad)
70        assert ws.ndim == stride.ndim and ws.ndim == pad.ndim
71        assert ws.ndim == 1
72        if ws.dtype not in theano.tensor.int_dtypes:
73            raise TypeError('Window shape parameters must be ints.')
74        if stride.dtype not in theano.tensor.int_dtypes:
75            raise TypeError('Stride parameters must be ints.')
76        if pad.dtype not in theano.tensor.int_dtypes:
77            raise TypeError('Padding parameters must be ints.')
78
79        ws = theano.tensor.cast(ws, 'int64')
80        stride = theano.tensor.cast(stride, 'int64')
81        pad = theano.tensor.cast(pad, 'int64')
82
83        return Apply(self, [inp, ws, stride, pad], [inp.type()])
84
85    def infer_shape(self, node, in_shapes):
86        ws, stride, pad = [node.inputs[1], node.inputs[2], node.inputs[3]]
87        shp = Pool.out_shape(in_shapes[0], ws, self.ignore_border, stride,
88                             pad, self.ndim)
89        return [shp]
90
91    def grad(self, inp, grads):
92        img, ws, stride, pad = inp
93        grad, = grads
94
95        grad = gpu_contiguous(grad)
96
97        disc = [theano.gradient.DisconnectedType()() for i in inp[1:]]
98        if self.mode == 'max':
99            out = self(img, ws, stride, pad)
100            g_out = GpuMaxPoolGrad(ndim=self.ndim,
101                                   ignore_border=self.ignore_border)(
102                                       img, out, grad, ws, stride, pad)
103            return [g_out] + disc
104        else:
105            g_out = GpuAveragePoolGrad(ndim=self.ndim,
106                                       ignore_border=self.ignore_border,
107                                       mode=self.mode)(
108                                           img, grad, ws, stride, pad)
109            return [g_out] + disc
110
111    def connection_pattern(self, node):
112        return [[1], [0], [0], [0]]
113
114    def R_op(self, inputs, eval_points):
115        if self.mode != 'max':
116            # Rop for average or sum is simply pooling evaluated at eval point
117            eval_inputs = [eval_points[0]] + inputs[1:]
118            return [self(*eval_inputs)]
119
120        # R_op can receive None as eval_points.
121        # That mean there is no diferientiable path through that input
122        # If this imply that you cannot compute some outputs,
123        # return None for those.
124        if eval_points[0] is None:
125            return [None]
126        z = self(*inputs)
127        x, ws, stride, pad = inputs
128        return [
129            GpuDownsampleFactorMaxGradGrad(self.ignore_border, self.mode,
130                                           self.ndim)(x, z, eval_points[0], ws,
131                                                      stride, pad)
132        ]
133
134
135class GpuMaxPoolGrad(CGpuKernelBase):
136    """
137    Implement the grad of max pooling on the gpu.
138
139    """
140    __props__ = ('ignore_border', 'mode', 'ndim')
141
142    def __init__(self, ignore_border, mode='max', ndim=2):
143        self.ndim = ndim
144        self.ignore_border = ignore_border
145        self.mode = mode
146        CGpuKernelBase.__init__(self, ['c_code/pool_max_grad.c'],
147                                'APPLY_SPECIFIC(max_pool_grad)')
148        assert mode == 'max'
149        assert ndim in [2, 3]
150
151    def c_headers(self):
152        return ['gpuarray_api.h', 'gpuarray_helper.h', 'numpy_compat.h']
153
154    def c_header_dirs(self):
155        return [gpuarray_helper_inc_dir(), pygpu.get_include()]
156
157    def make_node(self, inp, out, out_grad, ws, stride=None, pad=None):
158        ctx_name = infer_context_name(inp, out, out_grad)
159        nd = self.ndim
160        inp = as_gpuarray_variable(inp, ctx_name)
161        assert (inp.ndim == nd + 2)
162        out = as_gpuarray_variable(out, ctx_name)
163        assert (out.ndim == nd + 2)
164        out_grad = as_gpuarray_variable(out_grad, ctx_name)
165        assert (out_grad.ndim == nd + 2)
166
167        assert (out_grad.ndim == inp.ndim)
168        assert (inp.ndim == out.ndim)
169
170        if stride is None:
171            stride = ws
172        if pad is None:
173            pad = (0,) * nd
174        ws = as_tensor_variable(ws)
175        stride = as_tensor_variable(stride)
176        pad = as_tensor_variable(pad)
177        assert ws.ndim == stride.ndim and ws.ndim == pad.ndim
178        assert ws.ndim == 1
179        if ws.dtype not in theano.tensor.int_dtypes:
180            raise TypeError('Window shape parameters must be ints.')
181        if stride.dtype not in theano.tensor.int_dtypes:
182            raise TypeError('Stride parameters must be ints.')
183        if pad.dtype not in theano.tensor.int_dtypes:
184            raise TypeError('Padding parameters must be ints.')
185
186        ws = theano.tensor.cast(ws, 'int64')
187        stride = theano.tensor.cast(stride, 'int64')
188        pad = theano.tensor.cast(pad, 'int64')
189
190        return Apply(self, [inp, out, out_grad, ws, stride, pad], [inp.type()])
191
192    def infer_shape(self, node, in_shapes):
193        return [in_shapes[0]]
194
195    def grad(self, inp, grads):
196        x, maxout, gz, ws, stride, pad = inp
197        ggx, = grads
198        return ([theano.tensor.zeros_like(x),
199                 theano.tensor.zeros_like(maxout),
200                 GpuDownsampleFactorMaxGradGrad(ndim=self.ndim,
201                                                ignore_border=self.ignore_border)(
202                                                    x, maxout, ggx, ws, stride, pad)] +
203                [theano.tensor.DisconnectedType()() for i in inp[3:]])
204
205    def connection_pattern(self, node):
206        return [[1], [1], [1], [0], [0], [0]]
207
208
209class GpuAveragePoolGrad(CGpuKernelBase):
210    """
211    Implement the grad of average pooling on the gpu.
212
213    """
214    __props__ = ('ignore_border', 'mode', 'ndim')
215    params_type = ParamsType(mode=PoolingMode_t, context=gpu_context_type)
216
217    def __init__(self, ignore_border, mode='max', ndim=2):
218        self.ndim = ndim
219        self.ignore_border = ignore_border
220        if mode == 'average':
221            mode = 'average_inc_pad'
222        self.mode = mode
223        CGpuKernelBase.__init__(self, ['c_code/pool_ave_grad.c'],
224                                'APPLY_SPECIFIC(ave_pool_grad)')
225        assert mode in ('sum', 'average_inc_pad', 'average_exc_pad')
226        assert ndim in [2, 3]
227
228    def get_params(self, node):
229        return self.params_type.get_params(self, context=node.inputs[0].type.context)
230
231    def c_headers(self):
232        return ['gpuarray_api.h', 'gpuarray_helper.h', 'numpy_compat.h']
233
234    def c_header_dirs(self):
235        return [gpuarray_helper_inc_dir(), pygpu.get_include()]
236
237    def make_node(self, inp, out_grad, ws, stride=None, pad=None):
238        ctx_name = infer_context_name(inp, out_grad)
239        nd = self.ndim
240        inp = as_gpuarray_variable(inp, ctx_name)
241        assert (inp.ndim == nd + 2)
242        out_grad = as_gpuarray_variable(out_grad, ctx_name)
243        assert (out_grad.ndim == nd + 2)
244
245        assert (out_grad.ndim == inp.ndim)
246
247        if stride is None:
248            stride = ws
249        if pad is None:
250            pad = (0,) * nd
251        elif isinstance(pad, (tuple, list)):
252            if max(pad) != 0 and not self.mode == 'average_exc_pad':
253                raise ValueError('Padding must be zero for average_exc_pad')
254        ws = as_tensor_variable(ws)
255        stride = as_tensor_variable(stride)
256        pad = as_tensor_variable(pad)
257        assert ws.ndim == stride.ndim and ws.ndim == pad.ndim
258        assert ws.ndim == 1
259        if ws.dtype not in theano.tensor.int_dtypes:
260            raise TypeError('Window shape parameters must be ints.')
261        if stride.dtype not in theano.tensor.int_dtypes:
262            raise TypeError('Stride parameters must be ints.')
263        if pad.dtype not in theano.tensor.int_dtypes:
264            raise TypeError('Padding parameters must be ints.')
265
266        ws = theano.tensor.cast(ws, 'int64')
267        stride = theano.tensor.cast(stride, 'int64')
268        pad = theano.tensor.cast(pad, 'int64')
269
270        return Apply(self, [inp, out_grad, ws, stride, pad], [inp.type()])
271
272    def infer_shape(self, node, in_shapes):
273        return [in_shapes[0]]
274
275    def grad(self, inp, grads):
276        x, gz, ws, stride, pad = inp
277        ggx, = grads
278        return ([theano.tensor.zeros_like(x),
279                 GpuPool(ignore_border=self.ignore_border,
280                         ndim=self.ndim, mode=self.mode)(
281                             ggx, ws, stride, pad)] +
282                [theano.gradient.DisconnectedType()() for i in inp[2:]])
283
284    def connection_pattern(self, node):
285        return [[1], [1], [0], [0], [0]]
286
287
288class GpuDownsampleFactorMaxGradGrad(CGpuKernelBase):
289    """
290    Implement the grad of downsample with max on the gpu.
291
292    """
293    __props__ = ('ignore_border', 'mode', 'ndim')
294
295    def __init__(self, ignore_border, mode='max', ndim=2):
296        self.ndim = ndim
297        self.ignore_border = ignore_border
298        self.mode = mode
299        CGpuKernelBase.__init__(self, ['c_code/pool_grad_grad.c'],
300                                'APPLY_SPECIFIC(pool_grad_grad)')
301        assert self.mode == 'max'
302        assert self.ndim in [2, 3]
303
304    def c_headers(self):
305        return ['gpuarray_api.h', 'gpuarray_helper.h', 'numpy_compat.h']
306
307    def c_header_dirs(self):
308        return [gpuarray_helper_inc_dir(), pygpu.get_include()]
309
310    def make_node(self, inp, out, out_grad, ws, stride=None, pad=None):
311        ctx_name = infer_context_name(inp, out, out_grad)
312        nd = self.ndim
313        inp = as_gpuarray_variable(inp, ctx_name)
314        assert (inp.ndim == nd + 2)
315        out = as_gpuarray_variable(out, ctx_name)
316        assert (out_grad.ndim == nd + 2)
317        out_grad = as_gpuarray_variable(out_grad, ctx_name)
318        assert (out.ndim == nd + 2)
319
320        assert (out_grad.ndim == inp.ndim)
321        assert (inp.ndim == out.ndim)
322
323        if stride is None:
324            stride = ws
325        if pad is None:
326            pad = (0,) * nd
327        ws = as_tensor_variable(ws)
328        stride = as_tensor_variable(stride)
329        pad = as_tensor_variable(pad)
330        assert ws.ndim == stride.ndim and ws.ndim == pad.ndim
331        assert ws.ndim == 1
332        if ws.dtype not in theano.tensor.int_dtypes:
333            raise TypeError('Window shape parameters must be ints.')
334        if stride.dtype not in theano.tensor.int_dtypes:
335            raise TypeError('Stride parameters must be ints.')
336        if pad.dtype not in theano.tensor.int_dtypes:
337            raise TypeError('Padding parameters must be ints.')
338
339        ws = theano.tensor.cast(ws, 'int64')
340        stride = theano.tensor.cast(stride, 'int64')
341        pad = theano.tensor.cast(pad, 'int64')
342
343        return Apply(self, [inp, out, out_grad, ws, stride, pad], [inp.type()])
344
345    def infer_shape(self, node, in_shapes):
346        return [in_shapes[1]]
347
348    def grad(self, inp, grads):
349        x, maxout, ggx, ws, stride, pad = inp
350        gz, = grads
351        return ([theano.tensor.zeros_like(x),
352                theano.tensor.zeros_like(maxout),
353                GpuMaxPoolGrad(ignore_border=self.ignore_border,
354                               ndim=self.ndim)(
355                                   x, maxout, gz, ws, stride, pad)] +
356                [theano.gradient.DisconnectedType()() for i in inp[3:]])
357
358    def connection_pattern(self, node):
359        return [[1], [1], [1], [0], [0], [0]]
360
361
362class GpuMaxPoolRop(CGpuKernelBase):
363    """
364    Implements the R-operator for the downsample operation.
365
366    """
367    __props__ = ('ignore_border', 'mode', 'ndim')
368    params_type = ParamsType(ignore_border=bool_t, context=gpu_context_type)
369
370    def __init__(self, ignore_border, mode='max', ndim=2):
371        self.ndim = ndim
372        self.ignore_border = ignore_border
373        self.mode = mode
374        CGpuKernelBase.__init__(self, ['c_code/pool_max_rop.c'],
375                                'APPLY_SPECIFIC(max_pool_rop)')
376        assert mode == 'max'
377        assert ndim in [2, 3]
378
379    def get_params(self, node):
380        return self.params_type.get_params(self, context=node.inputs[0].type.context)
381
382    def c_headers(self):
383        return ['gpuarray_api.h', 'gpuarray_helper.h', 'numpy_compat.h']
384
385    def c_header_dirs(self):
386        return [gpuarray_helper_inc_dir(), pygpu.get_include()]
387
388    def make_node(self, inp, eval_point, ws, stride=None, pad=None):
389        ctx_name = infer_context_name(inp)
390        nd = self.ndim
391        inp = as_gpuarray_variable(inp, ctx_name)
392        assert (inp.ndim == nd + 2)
393        eval_point = as_gpuarray_variable(eval_point, ctx_name)
394        assert (eval_point.ndim == nd + 2)
395
396        if stride is None:
397            stride = ws
398        if pad is None:
399            pad = (0,) * nd
400        elif isinstance(pad, (tuple, list)):
401            if max(pad) != 0 and not self.ignore_border:
402                raise ValueError('Padding works only with ignore_border=True')
403            if isinstance(ws, (tuple, list)):
404                if any(pad[i] >= ws[i] for i in range(nd)):
405                    raise ValueError('Padding must be smaller than strides')
406
407        ws = as_tensor_variable(ws)
408        stride = as_tensor_variable(stride)
409        pad = as_tensor_variable(pad)
410        assert ws.ndim == stride.ndim and ws.ndim == pad.ndim
411        assert ws.ndim == 1
412        if ws.dtype not in theano.tensor.int_dtypes:
413            raise TypeError('Window shape parameters must be ints.')
414        if stride.dtype not in theano.tensor.int_dtypes:
415            raise TypeError('Stride parameters must be ints.')
416        if pad.dtype not in theano.tensor.int_dtypes:
417            raise TypeError('Padding parameters must be ints.')
418
419        ws = theano.tensor.cast(ws, 'int64')
420        stride = theano.tensor.cast(stride, 'int64')
421        pad = theano.tensor.cast(pad, 'int64')
422
423        return Apply(self, [inp, eval_point, ws, stride, pad], [eval_point.type()])
424
425    def infer_shape(self, node, in_shapes):
426        ws, stride, pad = [node.inputs[2], node.inputs[3], node.inputs[4]]
427        shp = Pool.out_shape(in_shapes[0], ws, self.ignore_border, stride,
428                             pad, self.ndim)
429        return [shp]
430