PyTorch デザインノート : PyTorch を拡張する (翻訳/解説)
翻訳 : (株)クラスキャット セールスインフォメーション
作成日時 : 05/26/2018 (0.4.0)
* 本ページは、PyTorch Doc Notes の – Extending PyTorch を動作確認・翻訳した上で適宜、補足説明したものです：
* サンプルコードの動作確認はしておりますが、適宜、追加改変している場合もあります。
* ご自由にリンクを張って頂いてかまいませんが、sales-info@classcat.com までご一報いただけると嬉しいです。
このノートでは torch.nn, torch.autograd を拡張して私達の C ライブラリを使用してカスタム C エクステンションを書く方法をカバーします。
torch.autograd を拡張する
autograd への演算の追加は各演算に対して新しい Function サブクラスの実装を必要とします。Function s は結果と勾配を計算するために autograd が使用するもので、演算履歴をエンコードします。総ての新しい関数は 2 つのメソッドを実装することを貴方に要求します :
- forward() – 演算を遂行するコードです。それは貴方が望むだけの引数を取ることができ、デフォルト値を指定すればそれらの幾つかはオプションを持ちます。ここでは総ての種類の Python オブジェクトが受け取れます。Variable 引数は呼び出し前に Tensor s に変換されてそれらの使用はグラフに登録されます。ロジックはリスト/辞書/任意の他のデータ構造をトラバースしません、そして呼び出しへの直接の引数である Variable を考慮するだけです。貴方は単一の Tensor 出力か、複数の出力があれば Tensor s の タプル を返すことができます。また、forword() からのみ呼び出すことができる有用なメソッドの完全な記述を見つけるためには Function の docs を参照してください。
- backward() – 勾配式 (= gradient formula) です。それは出力があるだけの数の Variable 引数が与えられ、それらの各々はその出力に関する勾配を表します。それは入力があるだけの数の Variable s を返すべきで、それらの各々はその対応する入力に関する勾配を含みます。もし貴方の入力が勾配を必要としない (needs_input_grad 参照) か、あるいは非-Variable オブジェクトである場合、None を返すことができます。また、forward() へのオプション引数を持つ場合は入力がある以上の勾配を返すことができます、それらが総て None である限り。
下で追加コメントを伴う、torch.nn からの Linear Function のためのコードを見つけることができます :
# Funtion から継承
class LinearFunction(Function):
# foward と backward の両者は @staticmethod であることに注意してください。
@staticmethod
# bias はオプション引数です。
def forward(ctx, input, weight, bias=None):
ctx.save_for_backward(input, weight, bias)
output = input.mm(weight.t())
if bias is not None:
output += bias.unsqueeze(0).expand_as(output)
return output
# この関数は単一の出力だけを持ちますので、1 つの勾配だけを取ります。
@staticmethod
def backward(ctx, grad_output):
# This is a pattern that is very convenient - at the top of backward
# unpack saved_tensors and initialize all gradients w.r.t. inputs to
# None. Thanks to the fact that additional trailing Nones are
# ignored, the return statement is simple even when the function has
# optional inputs.
input, weight, bias = ctx.saved_tensors
grad_input = grad_weight = grad_bias = None
# These needs_input_grad checks are optional and there only to
# improve efficiency. If you want to make your code simpler, you can
# skip them. Returning gradients for inputs that don't require it is
# not an error.
if ctx.needs_input_grad[0]:
grad_input = grad_output.mm(weight)
if ctx.needs_input_grad[1]:
grad_weight = grad_output.t().mm(input)
if bias is not None and ctx.needs_input_grad[2]:
grad_bias = grad_output.sum(0).squeeze(0)
return grad_input, grad_weight, grad_bias
今、これらのカスタム ops の使用をより簡単にするために、それらの apply メソッドをエイリアスすることを推奨します :
linear = LinearFunction.apply
ここで、非-Variable 引数によりパラメータ化された追加の関数サンプルを与えます :
class MulConstant(Function):
@staticmethod
def forward(ctx, tensor, constant):
# ctx is a context object that can be used to stash information
# for backward computation
ctx.constant = constant
return tensor * constant
@staticmethod
def backward(ctx, grad_output):
# We return as many input gradients as there were arguments.
# Gradients of non-Tensor arguments to forward must be None.
return grad_output * ctx.constant, None
貴方は多分、実装した backward メソッドが実際に貴方の関数の導関数を計算するかを確認することを望むでしょう。それは小さい有限の差異を使用して数値近似を比較することにより可能です :
from torch.autograd import gradcheck # gradcheck takes a tuple of tensors as input, check if your gradient # evaluated with these tensors are close enough to numerical # approximations and returns True if they all verify this condition. input = (Variable(torch.randn(20,20).double(), requires_grad=True), Variable(torch.randn(30,20).double(), requires_grad=True),) test = gradcheck(Linear.apply, input, eps=1e-6, atol=1e-4) print(test)
torch.nn を拡張する
nn は 2 つの種類のインターフェイスを持ちます – モジュールとそれらの functional バージョンです。それを両者の方法で拡張できますが、(任意のパラメータやバッファを保持する) 総ての種類の層のためにはモジュールを使用することを推奨し、そして活性化関数、プーリング, etc. のようなパラメータ-less 演算のためには functional 形式の使用を勧めます。
演算の functional バージョンの追加は上のセクションで既に完全にカバーされています。
Module の追加
nn は autograd を大いに利用しますので、新しい Module の追加は演算を遂行して勾配を計算可能な Function の実行を必要とします。ここから私達は Linear モジュールを実装することを望みそして上のリスティング内のように実装された関数を持つことを仮定しましょう。これを追加するために必要な非常に少ないコードがあります。今、実装される必要がある 2 つの関数があります :
- __init__ (オプション) – カーネルサイズ、特徴数, etc. のような引数を取ってパラメータとバッファを初期化します。
- forward() – Function をインスタンス化してそれを演算を遂行するために使用します。それは上で示された functional ラッパーに非常に類似しています。
これはどのように Linear モジュールが実装できるかです :
class Linear(nn.Module):
def __init__(self, input_features, output_features, bias=True):
super(Linear, self).__init__()
self.input_features = input_features
self.output_features = output_features
# nn.Parameter is a special kind of Variable, that will get
# automatically registered as Module's parameter once it's assigned
# as an attribute. Parameters and buffers need to be registered, or
# they won't appear in .parameters() (doesn't apply to buffers), and
# won't be converted when e.g. .cuda() is called. You can use
# .register_buffer() to register buffers.
# nn.Parameters require gradients by default.
self.weight = nn.Parameter(torch.Tensor(output_features, input_features))
if bias:
self.bias = nn.Parameter(torch.Tensor(output_features))
else:
# You should always register all possible parameters, but the
# optional ones can be None if you want.
self.register_parameter('bias', None)
# Not a very smart way to initialize weights
self.weight.data.uniform_(-0.1, 0.1)
if bias is not None:
self.bias.data.uniform_(-0.1, 0.1)
def forward(self, input):
# See the autograd section for explanation of what happens here.
return LinearFunction.apply(input, self.weight, self.bias)
def extra_repr(self):
# (Optional)Set the extra information about this module. You can test
# it by printing an object of this class.
return 'in_features={}, out_features={}, bias={}'.format(
self.in_features, self.out_features, self.bias is not None
カスタム C エクステンションを書く
Coming soon. 当面はサンプルを GitHub で見つけることができます。
以上