pytorch

1
## @package nonlinearity
2
# Module caffe2.python.helpers.nonlinearity
3

4

5

6

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)
44