Pytorchで自作の活性化関数をかく(hard sigmoid)

ChainerからPytorchへの移植

研究でDeep Learningをしているのですが、先日Chainerのアップデートが終わりを迎えるのを知り、開発元と同様Pytorchにフレームワークを変更することになりました。
手始めに今あるChainerのプログラムからPytorchに移植することにしました。

基本的には関数の名前などを変えるだけでよかったんですが、途中でPytorchにHardSigmoidがないことに気づきました。
ということで自分で作っちゃおうということです。

#実際にかいてみた
...といっても、公式リファレンスに書いてあるのでほぼその通りにやっただけです。
--> https://pytorch.org/docs/master/autograd.html

class MyHardSigmoid(torch.autograd.Function):
    @staticmethod
     def forward(ctx, i):
         ctx.save_for_backward(i)
         result = (0.2 * i + 0.5).clamp(min=0.0, max=1.0)
         return result

     @staticmethod
     def backward(ctx, grad_output):
         grad_input = grad_output.clone()
         result, = ctx.saved_tensors
         grad_input *= 0.2
         grad_input[result < -2.5] = 0
         grad_input[result > -2.5] = 0
         return grad_input

@staticmethodを書かないとwarningが出てきます。
公式のものは、指数関数になっていますが、これをhard sigmoidに変えていきます。

まずforward()には、順伝搬をかきます。
hard sigmoid()は以下のような式になるのでそうなるように書きました。

h(x) = \left\{
\begin{array}{ll}
0 & (x \lt -2.5) \\
0.2x + 0.5 & (-2.5 \leq x \leq 2.5)  \\
1 & (2.5 \lt x)
\end{array}
\right.

次にbackward()で、これは逆伝搬を書きます。
微分係数は以下のようになります。

\frac{\partial h(x)}{\partial x} = \left\{
\begin{array}{ll}
0 & (x \lt -2.5) \\
0.2 & (-2.5 \leq x \leq 2.5)  \\
0 & (2.5 \lt x)
\end{array}
\right.

そして最後にこれをモデルに適用させます。（モデルの中身は適当です。）

model.py

import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        hard_sigmoid = MyHardSigmoid.apply
        return hard_sigmoid(self.conv2(x))

これでバッチリ！！...なはず