PyTorch 1.3 Tutorials : 名前付き tensor (試験的) : PyTorch の名前付き tensor へのイントロダクション
翻訳 : (株)クラスキャット セールスインフォメーション
作成日時 : 01/16/2020 (1.3.1)

* 本ページは、PyTorch 1.3 Tutorials の以下のページを翻訳した上で適宜、補足説明したものです：

• Text : Introduction to Named Tensors in PyTorch

* サンプルコードの動作確認はしておりますが、必要な場合には適宜、追加改変しています。
* ご自由にリンクを張って頂いてかまいませんが、sales-info@classcat.com までご一報いただけると嬉しいです。

 

名前付き tensor (試験的) : PyTorch の名前付き tensor へのイントロダクション

名前付き tensor はユーザに明示的な名前を tensor 次元と関連付けることを可能にすることにより tensor の利用をより容易にすることを目標にしています。殆どの場合、次元パラメータを取る演算は、位置により次元を追跡する必要性を回避するために、次元名を受け取ります。加えて、特別な安全性を提供するため、名前付き tensor は API が実行時に正しく使われているかを自動的に確認するために名前を使用します。名前はまた次元を再配置するためにも使用できます、例えば、「位置によりブロードキャスト」するよりも「名前でブロードキャスト」することをサポートするためにです。

このチュートリアルは 1.3 launch で含まれる機能へのガイドとして意図されています。その終わりまでには、次ができるようになります :

• 名前付き次元で tensor を作成し、またそれらの次元を除去または名前変更する
• 演算が次元名をどのように伝播するかの基本を理解する
• 次元を名前付けることが 2 つの主要な領域でより明瞭なコードを可能にするかを見る :
  • 演算をブロードキャストする
  • 次元を平坦化と非平坦化する (= flattening と unflattening)

最後に、名前付き tensor を使用して multi-head attention モジュールを書くことによりこれを実践します。

PyTorch の名前付き tensor は Sasha Rush との協力でインスパイアされて成されました。

Sasha は元のアイデアと概念の証拠 (= proof) を 2019 年 1 月のブログ投稿 で提案しました。

 

基本: 名前付き次元

PyTorch は今では Tensor に名前付き次元を持つことを可能にします ; ファクトリ関数は新しい names 引数を取り、これは名前を各次元と関連付けます。これは次のような殆どのファクトリ関数で動作します :

• tensor
• empty
• ones
• zeros
• randn
• rand

ここで names を伴う tensor を構築します :

import torch
imgs = torch.randn(1, 2, 2, 3, names=('N', 'C', 'H', 'W'))
print(imgs.names)

('N', 'C', 'H', 'W')

元の名前付き tensor ブログ投稿 とは違い、名前付き次元は順序付けされます : tensor.names[i] は tensor の i th 次元の名前です。

Tensor の次元を名前変更するには 2 つの方法があります :

# Method #1: set the .names attribute (this changes name in-place)
imgs.names = ['batch', 'channel', 'width', 'height']
print(imgs.names)

# Method #2: specify new names (this changes names out-of-place)
imgs = imgs.rename(channel='C', width='W', height='H')
print(imgs.names)

('batch', 'channel', 'width', 'height')
('batch', 'C', 'W', 'H')

names を除去する好ましい方法は tensor.rename(None) を呼び出すことです :

imgs = imgs.rename(None)
print(imgs.names)

(None, None, None, None)

非名前付き tensor は依然として通常のように動作してそれらの repr に names を持ちません。

unnamed = torch.randn(2, 1, 3)
print(unnamed)
print(unnamed.names)

tensor([[[ 0.5911, -0.9557,  0.0103]],

        [[-0.6704, -0.7485,  0.8281]]])
(None, None, None)

名前付き tensor は総ての次元が名前付けられることを要求しません。

imgs = torch.randn(3, 1, 1, 2, names=('N', None, None, None))
print(imgs.names)

('N', None, None, None)

名前付き tensor は非名前付き tensor と共存できますので、名前付きと非名前付き tensor の両者で動作する名前付き tensor-aware コードを書く良い方法が必要です。次元を再調整 (= refine) して非名前付き dim を名前付き dim にリフトするために tensor.refine_names(*names) を使用します。次元の再調整は次の制約を持つ “rename” として定義されます :

• None dim は任意の名前を持つように再調整できます。
• 名前付き dim は同じ名前を持つように再調整できるだけです。

imgs = torch.randn(3, 1, 1, 2)
named_imgs = imgs.refine_names('N', 'C', 'H', 'W')
print(named_imgs.names)

# Refine the last two dims to 'H' and 'W'. In Python 2, use the string '...'
# instead of ...
named_imgs = imgs.refine_names(..., 'H', 'W')
print(named_imgs.names)


def catch_error(fn):
    try:
        fn()
        assert False
    except RuntimeError as err:
        err = str(err)
        if len(err) > 180:
            err = err[:180] + "..."
        print(err)


named_imgs = imgs.refine_names('N', 'C', 'H', 'W')

# Tried to refine an existing name to a different name
catch_error(lambda: named_imgs.refine_names('N', 'C', 'H', 'width'))

('N', 'C', 'H', 'W')
(None, None, 'H', 'W')
refine_names: cannot coerce Tensor['N', 'C', 'H', 'W'] to Tensor['N', 'C', 'H', 'width'] because 'W' is different from 'width' at index 3

殆どの単純な演算は名前を伝播します。名前付き tensor の最終的な目標は総ての演算について合理的で、直感的な流儀で名前を伝播することです。多くの一般的な演算のためのサポートは 1.3 リリースの時点で追加されました ; ここでは、例として、.abs() :

print(named_imgs.abs().names)

('N', 'C', 'H', 'W')

 

アクセサとリダクション

次元を参照するために位置次元の代わりに次元名を使用できます。これらの演算はまた名前を伝播します。インデキシング (基本と上級) はまだ実装されていませんがロードマップにあります。上からの named_imgs tensor を使用して、以下を行なうことができます :

output = named_imgs.sum('C')  # Perform a sum over the channel dimension
print(output.names)

img0 = named_imgs.select('N', 0)  # get one image
print(img0.names)

('N', 'H', 'W')
('C', 'H', 'W')

 

名前推論

名前は名前推論と呼ばれる 2 ステップ過程の演算上で伝播されます。

1. 名前を確認する: 演算子は、ある次元名が一致しなければならないことを確認する自動チェックを実行時に遂行するかもしれません。
2. 名前を伝播する: 名前推論は出力名を出力 tensor に伝播します。

ブロードキャストなしの 2 つの 1-dim tensor を可算する非常に小さいサンプルを調べましょう。

x = torch.randn(3, names=('X',))
y = torch.randn(3)
z = torch.randn(3, names=('Z',))

名前を確認する: 最初に、これら 2 つの tensor の名前が一致するかを確認します。2 つの名前はそれらが等しい (文字列等値) か少なくとも一つが None である (None は基本的には特別なワイルドカード名) 場合に限り一致します。エラーになるこれら 3 つの唯一のものは、従って、 x + z です :

catch_error(lambda: x + z)

Error when attempting to broadcast dims ['X'] and dims ['Z']: dim 'X' and dim 'Z' are at the same position from the right but do not match.

名前を伝播する: 2 つの最も調整された名前を返すことにより 2 つの名前を統一します。x + y では、X は None よりもより調整されています。

print((x + y).names)

('X',)

殆どの名前推論ルールは簡単ですが、それらの幾つかは予期せぬセマンティクスを持つ可能性があります。遭遇しがちな 2 つを調べましょう : ブロードキャストと行列乗算です。

 

ブロードキャスト

名前付き tensor はブロードキャストの動作を変更しません ; それらは依然として位置によりブロードキャストします。けれども、2 つの次元をそれらがブロードキャストできるか否かのために確認するとき、PyTorch はまたそれらの次元の名前が一致するかを確認します。

これは名前付き tensor がブロードキャストする演算の間に意図しない配置を防ぐ結果になります。下のサンプルでは、imgs に per_batch_scale を適用します。

imgs = torch.randn(2, 2, 2, 2, names=('N', 'C', 'H', 'W'))
per_batch_scale = torch.rand(2, names=('N',))
catch_error(lambda: imgs * per_batch_scale)

Error when attempting to broadcast dims ['N', 'C', 'H', 'W'] and dims ['N']: dim 'W' and dim 'N' are at the same position from the right but do not match.

名前なしでは、per_batch_scale tensor は imgs の最後の次元で配置されます、これえは私達が意図したものではありません。実際には per_batch_scale を imgs のバッチ次元で配置することにより演算を遂行することを望みました。下でカバーされる、tensor を名前でどのように配置するかについては新しい「名前による明示的なブロードキャスト」機能を見てください。

 

行列乗算

torch.mm(A, B) は A の 2 番目の dim と B の最初の dim の間のドット積を遂行し、A の最初の dim と B の 2 番目の dim を持つ tensor を返します (torch.matmul, torch.mv と torch.dot のような他の matmul 関数は同様に動作します)。

markov_states = torch.randn(128, 5, names=('batch', 'D'))
transition_matrix = torch.randn(5, 5, names=('in', 'out'))

# Apply one transition
new_state = markov_states @ transition_matrix
print(new_state.names)

('batch', 'out')

見れるように、行列乗算は縮小される次元が同じ名前を持つかどうかを確認しません。

次に、名前付き tensor が可能にする 2 つの新しい動作をカバーします : 名前による明示的なブロードキャストと名前による次元の平坦化と非平坦化です。

 

新しい動作 : 名前による明示的なブロードキャスト

マルチ次元で作業することについての主要な苦情の一つは演算ができるように「ダミー」次元を unsqueeze する必要性です。例えば、前の per-batch-scale サンプルで、非名前付き tensor では次のようにするでしょう :

imgs = torch.randn(2, 2, 2, 2)  # N, C, H, W
per_batch_scale = torch.rand(2)  # N

correct_result = imgs * per_batch_scale.view(2, 1, 1, 1)  # N, C, H, W
incorrect_result = imgs * per_batch_scale.expand_as(imgs)
assert not torch.allclose(correct_result, incorrect_result)

名前を使うことによりこれらの演算をより安全にすることができます (そして容易に次元の数に不可知です)。適切なところで one-sized 次元を追加して (tensor.align_to(*names) もまた動作します)、other.names で指定される順序に適合するために tensor の次元を並べ替える (= permute) 新しい tensor.align_as(other) 演算を提供します :

imgs = imgs.refine_names('N', 'C', 'H', 'W')
per_batch_scale = per_batch_scale.refine_names('N')

named_result = imgs * per_batch_scale.align_as(imgs)
# note: named tensors do not yet work with allclose
assert torch.allclose(named_result.rename(None), correct_result)

 

新しい動作 : 名前による次元の平坦化と非平坦化

一つの一般的な演算は次元の平坦化と非平坦化です。ユーザは view, reshape または flatten のいずれかを使用してこれを今すぐ、遂行できます ; ユースケースは tensor を (ある数の次元 (i.e., conv2d は 4D 入力を取ります) を持つ入力を取らなければならない) 演算子に送るためにバッチ次元を平坦化することを含みます。

これらの演算を view や reshape よりもセマンティクス的に意味があるものにするため、新しい tensor.unflatten(dim, namedshape) メソッドを導入して名前とともに動作するように flatten を更新します : tensor.flatten(dims, new_dim)。

flatten は隣接する (= adjacent) 次元を平坦化できるだけですが不連続な (= non-contiguou) dim 上でも動作します。dim をどのように unflatten するかを指定するために unflatten に 名前付き shape を渡さなければなりません、これは (dim size) タプルのリストです。unflatten のために flatten の間にサイズをセーブすることは可能ですがそれをまだしていません。

imgs = imgs.flatten(['C', 'H', 'W'], 'features')
print(imgs.names)

imgs = imgs.unflatten('features', (('C', 2), ('H', 2), ('W', 2)))
print(imgs.names)

('N', 'features')
('N', 'C', 'H', 'W')

 

Autograd サポート

Autograd は現在総ての tensor 上の名前を無視してそれらを標準的な tensor のように扱います。勾配計算は正しいですが名前が与えてくれる安全性は失います。名前の処理の autograd への導入はロードマップ上です。

x = torch.randn(3, names=('D',))
weight = torch.randn(3, names=('D',), requires_grad=True)
loss = (x - weight).abs()
grad_loss = torch.randn(3)
loss.backward(grad_loss)

correct_grad = weight.grad.clone()
print(correct_grad)  # Unnamed for now. Will be named in the future

weight.grad.zero_()
grad_loss = grad_loss.refine_names('C')
loss = (x - weight).abs()
# Ideally we'd check that the names of loss and grad_loss match, but we don't
# yet
loss.backward(grad_loss)

print(weight.grad)  # still unnamed
assert torch.allclose(weight.grad, correct_grad)

tensor([0.5111, 2.3763, 0.0755])
tensor([0.5111, 2.3763, 0.0755])

 

他のサポートされる (そして非サポートの) 特徴

1.3 リリースでサポートされるものの詳細な分析については こちらを見てください。

特に、現在サポートされていない 3 つの重要な特徴を呼び出すことを望みます :

• 名前付き tensor を torch.save または torch.load を通してセーブまたはロードします
• torch.multiprocessing 経由のマルチ処理
• JIT サポート ; 例えば、次はエラーになります :

imgs_named = torch.randn(1, 2, 2, 3, names=('N', 'C', 'H', 'W'))


@torch.jit.script
def fn(x):
    return x


catch_error(lambda: fn(imgs_named))

NYI: Named tensors are currently unsupported in TorchScript. As a  workaround please drop names via `tensor = tensor.rename(None)`. (guardAgainstNamedTensor at /pytorch/torch/csrc/...

回避方法として、名前付き tensor をまだサポートしていない何かを使用する前に tensor = tensor.rename(None) で名前を破棄してください。

 

より長いサンプル : Multi-head attention

今は一般的な PyTorch nn.Module : multi-head attention を実装する完全なサンプルを調べます。読者は既に multi-head attention に馴染みがあることを仮定しています、この説明 か この説明 を確認してください。

ParlAI からの multi-head attention の実装を適応させます ; 特に こちら です。そのサンプルのコードを読み通してください ; それから、下のコードと比較します、(I), (II), (III) と (IV) がラベル付けられた 4 つの箇所があることに注意してください、そこでは名前付き tensor の使用がより可読なコードを可能にしています ; コードブロックの後でこれらの各々に飛び込みます。

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


class MultiHeadAttention(nn.Module):
    def __init__(self, n_heads, dim, dropout=0):
        super(MultiHeadAttention, self).__init__()
        self.n_heads = n_heads
        self.dim = dim

        self.attn_dropout = nn.Dropout(p=dropout)
        self.q_lin = nn.Linear(dim, dim)
        self.k_lin = nn.Linear(dim, dim)
        self.v_lin = nn.Linear(dim, dim)
        nn.init.xavier_normal_(self.q_lin.weight)
        nn.init.xavier_normal_(self.k_lin.weight)
        nn.init.xavier_normal_(self.v_lin.weight)
        self.out_lin = nn.Linear(dim, dim)
        nn.init.xavier_normal_(self.out_lin.weight)

    def forward(self, query, key=None, value=None, mask=None):
        # (I)
        query = query.refine_names(..., 'T', 'D')
        self_attn = key is None and value is None
        if self_attn:
            mask = mask.refine_names(..., 'T')
        else:
            mask = mask.refine_names(..., 'T', 'T_key')  # enc attn

        dim = query.size('D')
        assert dim == self.dim, \
            f'Dimensions do not match: {dim} query vs {self.dim} configured'
        assert mask is not None, 'Mask is None, please specify a mask'
        n_heads = self.n_heads
        dim_per_head = dim // n_heads
        scale = math.sqrt(dim_per_head)

        # (II)
        def prepare_head(tensor):
            tensor = tensor.refine_names(..., 'T', 'D')
            return (tensor.unflatten('D', [('H', n_heads), ('D_head', dim_per_head)])
                          .align_to(..., 'H', 'T', 'D_head'))

        assert value is None
        if self_attn:
            key = value = query
        elif value is None:
            # key and value are the same, but query differs
            key = key.refine_names(..., 'T', 'D')
            value = key
        dim = key.size('D')

        # Distinguish between query_len (T) and key_len (T_key) dims.
        k = prepare_head(self.k_lin(key)).rename(T='T_key')
        v = prepare_head(self.v_lin(value)).rename(T='T_key')
        q = prepare_head(self.q_lin(query))

        dot_prod = q.div_(scale).matmul(k.align_to(..., 'D_head', 'T_key'))
        dot_prod.refine_names(..., 'H', 'T', 'T_key')  # just a check

        # (III)
        attn_mask = (mask == 0).align_as(dot_prod)
        dot_prod.masked_fill_(attn_mask, -float(1e20))

        attn_weights = self.attn_dropout(F.softmax(dot_prod / scale,
                                                   dim='T_key'))

        # (IV)
        attentioned = (
            attn_weights.matmul(v).refine_names(..., 'H', 'T', 'D_head')
            .align_to(..., 'T', 'H', 'D_head')
            .flatten(['H', 'D_head'], 'D')
        )

        return self.out_lin(attentioned).refine_names(..., 'T', 'D')

 
(I) 入力 tensor dim を再調整する

def forward(self, query, key=None, value=None, mask=None):
    # (I)
    query = query.refine_names(..., 'T', 'D')

query = query.refine_names(…, ‘T’, ‘D’) は実行可能なドキュメントとして役立ちそして入力次元を名前付けられるように持ち上げます。それは最後の 2 次元が [‘T’, ‘D’] に再調整できて、後で行をダウンさせる潜在的にサイレントであるか混乱させるサイズ不適合エラーを回避します。

 
(II) prepare_head で次元を操作する

# (II)
def prepare_head(tensor):
    tensor = tensor.refine_names(..., 'T', 'D')
    return (tensor.unflatten('D', [('H', n_heads), ('D_head', dim_per_head)])
                  .align_to(..., 'H', 'T', 'D_head'))

注意すべき最初のことはコードが入力と出力次元をどのように明白に述べるかです : 入力 tensor は T と D dim で終わりそして出力 tensor は H, T と D_head dim で終わらなければなりません。

注意すべき 2 番目のことはコードが何が起きているかをどのように明白に記述するかです。prepare_head は key, query と value を取りそして埋め込み dim をマルチヘッドに分割し、最後に dim 順序を […, ‘H’, ‘T’, ‘D_head’] であるように再配置します。ParlAI は view と transpose 演算を使用して、prepare_head を次のように実装しています :

def prepare_head(tensor):
    # input is [batch_size, seq_len, n_heads * dim_per_head]
    # output is [batch_size * n_heads, seq_len, dim_per_head]
    batch_size, seq_len, _ = tensor.size()
    tensor = tensor.view(batch_size, tensor.size(1), n_heads, dim_per_head)
    tensor = (
        tensor.transpose(1, 2)
        .contiguous()
        .view(batch_size * n_heads, seq_len, dim_per_head)
    )
    return tensor

名前付き tensor 変種 (= variant) はより冗長 (= verbose) ですが、view と transpose よりセマンティックな意味を持つ ops を使用し、そして名前の形式で実行可能なドキュメントを含みます。

 
(III) 名前による明示的なブロードキャスト

def ignore():
    # (III)
    attn_mask = (mask == 0).align_as(dot_prod)
    dot_prod.masked_fill_(attn_mask, -float(1e20))

mask は通常は dim [N, T] (self attention の場合) か [N, T, T_key] (エンコーダ attention の場合) を持ち、一方で dot_prod は diim [N, H, T, T_key] を持ちます。mask を dot_prod と共に正しくブロードキャストさせるために、通常は self attention の場合に dim 1 と -1 を unsqueeze するか、エンコーダ attention の場合に dim 1 を unsqueeze するでしょう。 名前付き tensor を使用すれば、align_as を使用して単純に attn_mask を dot_prod に整列させます、そしてどこに dim を unsqueeze するかについて心配することを止めます。

 
(IV) align_to と flatten を使用してより多くの次元操作

def ignore():
    # (IV)
    attentioned = (
        attn_weights.matmul(v).refine_names(..., 'H', 'T', 'D_head')
        .align_to(..., 'T', 'H', 'D_head')
        .flatten(['H', 'D_head'], 'D')
    )

ここで、(II) 内のように、align_to and flatten は (より冗長であるにもかかわらず) view と transpose よりもセマンティクス的に意味があります。

 

サンプルを実行する

n, t, d, h = 7, 5, 2 * 3, 3
query = torch.randn(n, t, d, names=('N', 'T', 'D'))
mask = torch.ones(n, t, names=('N', 'T'))
attn = MultiHeadAttention(h, d)
output = attn(query, mask=mask)
# works as expected!
print(output.names)

('N', 'T', 'D')

上は期待したように動作します。更に、コードでバッチ次元の名前に全く言及していないことに注意してください。実際に、MultiHeadAttention モジュールはバッチ次元の存在に不可知です。

query = torch.randn(t, d, names=('T', 'D'))
mask = torch.ones(t, names=('T',))
output = attn(query, mask=mask)
print(output.names)

('T', 'D')

 

終わりに

Thank you for reading! 名前付き tensor は依然として開発の真っ最中です ; 改良のためのフィードバック and/or 提案を持つのであれば、issue を作成して私達に知らせてください。

 
以上