Caffe2 - Python API
A deep learning, cross platform ML framework
nonlinearity.py
1 ## @package nonlinearity
2 # Module caffe2.python.helpers.nonlinearity
3 from __future__ import absolute_import
4 from __future__ import division
5 from __future__ import print_function
6 from __future__ import unicode_literals
7 
8 from caffe2.python import core
9 
10 
11 def prelu(model, blob_in, blob_out, num_channels=1, slope_init=None,
12  **kwargs):
13  """PRelu"""
14  slope_init = (
15  slope_init if slope_init else ('ConstantFill', {'value': 0.25}))
16  if model.init_params:
17  slope = model.param_init_net.__getattr__(slope_init[0])(
18  [],
19  blob_out + '_slope',
20  shape=[num_channels],
21  **slope_init[1]
22  )
23  else:
24  slope = core.ScopedBlobReference(
25  blob_out + '_slope', model.param_init_net)
26 
27  model.AddParameter(slope)
28 
29  return model.net.PRelu([blob_in, slope], [blob_out])
30 
31 
32 def relu(model, blob_in, blob_out, use_cudnn=False, order="NCHW", **kwargs):
33  """Relu."""
34  if use_cudnn:
35  kwargs['engine'] = 'CUDNN'
36  return model.net.Relu(blob_in, blob_out, order=order, **kwargs)
37 
38 
39 def tanh(model, blob_in, blob_out, use_cudnn=False, order="NCHW", **kwargs):
40  """Tanh."""
41  if use_cudnn:
42  kwargs['engine'] = 'CUDNN'
43  return model.net.Tanh(blob_in, blob_out, order=order, **kwargs)