NumPyやPyTorchで使える超便利ツールを作った

この記事は，CAMPHOR- Advent Calendar 2018の7日目の記事です．
この記事は，NumPyやPyTorchなどの開発に使える「shape_commentator」という便利ツールを使った話です．このツールはpipでインストールできるので，サクッと試してみたい方はGitHubのページに飛んで使ってみてください．

加筆修正(2019/12/12) IPython拡張を追加したことでIPythonやJupyterNotebookでの使い心地が良くなったので，IPythonの例の部分を中心に一部修正しました．

はじめに

こんにちは．シバニャンです．最近は卒論で画像系の研究をしています．

卒業研究ではNumPyやPyTorchを使ってコーディングしているんですが，今までUnityのコードをC#で書いてVisual Studioのコード補完のぬるま湯に浸かってきた僕とっては，型が明示的でないコードで計算を書いていくのが難しく思えました．

NumPyの辛さ

具体的にNumPyで一番難しく感じた部分は，ソースコードを見ただけではテンソルのShapeが分からないことです．

(NumPy使ってるよ，面倒だよねという方はこの章は飛ばしてください．)
下のコードのように3つテンソルを作ってみましょう．aは2次元のベクトルで，bは縦(axis0)が2行，横(axis1)が3列の2x3行列です．
このとき，aのShapeは(2,)で，bのShapeは(2,3)というタプルで表されます．簡単ですね．
この調子でcは，axis0が2，axis1が3，axis2が4の大きさを持つテンソルになります．このShapeは(2,3,4)と表されます．

import numpy as np
a = np.array([2,3])
b = np.array([[1,2,3],[3,4,5]])
c = np.array([
    [[1,2,3,4],[2,3,4,5],[3,4,5,6]],
    [[2,3,4,5],[3,4,5,6],[4,5,6,7]]
])

print("====== a =====")
print("a: {}".format(a))
print("a.shape: {}".format(a.shape))

print("====== b =====")
print("b: {}".format(b))
print("b.shape: {}".format(b.shape))

print("====== c =====")
print("c: {}".format(c))
print("c.shape: {}".format(c.shape))

====== a =====
a: [2 3]
a.shape: (2,)
====== b =====
b: [[1 2 3]
 [3 4 5]]
b.shape: (2, 3)
====== c =====
c: [[[1 2 3 4]
  [2 3 4 5]
  [3 4 5 6]]

 [[2 3 4 5]
  [3 4 5 6]
  [4 5 6 7]]]
c.shape: (2, 3, 4)

なーんだ，こんなの見たらわかるじゃんと思った方，こっちならどうでしょう?
変数ab_h, ab_v, ab, AA, BBのShapeは一瞬で分かるでしょうか?

import numpy as np
a = np.array([1,2,3,4,5,6])
b = np.array([0,1,2,3,4,5])

ab_h = np.hstack((a,b))
ab_v = np.vstack((a,b))
ab = np.dot(a,b)
AA, BB = np.meshgrid(a,b)

文句言わずドキュメント読めよ，と思った方がいると思いますが，この例は分かりやすいNumPyのメソッドだから良いものの，これがもし昔に作ったメソッドや研究室の秘伝のソースだったら一瞬で分かるはずはありません．

shapeを知りたいだけだった場合にも，毎回ドキュメントを参照しながらソースをじっくり読み込んだり，毎回print文を書いてshapeを調べながら，本当は研究に使うべき精神力を消耗していったりすることが想像できると思います． (shapeだけでいいから知りたいという状況は，ニューラルネットワークのコードを書くときに入力のノード数と出力のノード数を合わせるときに必ず発生します．)

こういったことに労力を割かれているのは，NumPyなどで行列演算をしている人にとっては日常茶飯事で，もしかしたらこの面倒なプロセスに慣れてしまったかもしれません．

型・・・?

Pythonには型ヒントという機能がありますが，NumPyのShapeまでは(たぶん)サポートされていません．

もしあったとしても，自分で型ヒントを書きながら開発を進めるのはそれは結構なことですが，型ヒントの書いていないライブラリや研究室の秘伝のソースを利用しようとなると，結局print文でshapeを調べる羽目になるのです．僕はそんな面倒なことはしたくありません．

今回の僕の目的は，NumPyなどのShapeを一瞬で知りたいということです．あくまで実行時のShapeが分かれば数値計算を進めるのには十分なのです．

便利ツールを作った

NumPyなどでShapeが分からないという辛い問題を，僕はソースコードにコメントをつける「shape_commentator」というツールを作ることで解決しました．「NumPyのshapeをコメントしてくれる君」的なネーミングですが，可愛くないですか?

shape_commentatorはpipでインストールできます．

pip install shape_commentator

IPython, Jupyter Notebookで使う

まずはIPythonで使ってみましょう．shape_commentatorをインポートして，先ほどのコードのセルの上に

%%shape

というマジックコマンドを実行します．

(見た方が早いと思うので，IPython上での実行結果を貼ります．)

In [1]: import shape_commentator

In [2]: %%shape
   ...: import numpy as np
   ...: a = np.array([1,2,3,4,5,6])
   ...: b = np.array([0,1,2,3,4,5])
   ...: ab_h = np.hstack((a,b))
   ...: ab_v = np.vstack((a,b))
   ...: ab = np.dot(a,b)
   ...: AA, BB = np.meshgrid(a,b)
   ...: 
import numpy as np
a = np.array([1,2,3,4,5,6])  #_ (6,)
b = np.array([0,1,2,3,4,5])  #_ (6,)
ab_h = np.hstack((a,b))  #_ (12,)
ab_v = np.vstack((a,b))  #_ (2, 6)
ab = np.dot(a,b)  #_ ()
AA, BB = np.meshgrid(a,b)  #_ [(6, 6),(6, 6),]

すると，あらびっくり，Shapeのコメント付きのソースコードが手に入ります．これ，めっちゃ便利じゃないですか!?
うんうん，便利でしょう．公開したソースコードに丁寧にshapeのコメントが書いてあったら見る側も嬉しいと思います．

これによってshape_commentatorは，渡されたソースコードを実際に実行し，代入文を乗っ取ってshapeの情報を取り出してソースコードに付け加えて出力します． In[2]の最後の行のnp.meshgridは返り値がlistなんですが，listやtupleを展開してくれる機能をつけたので，いい感じに展開してくれます．

同様に，Jupyter Notebookでも同じように利用できます．(やり方は全く同じなので省略)

ファイルに対して使う

IPythonではなくて，Pythonのファイルに対してshape_commentatorを実行することもできます．

src.pyというファイルに対して実行するのであれば，

import numpy as np
a = np.array([1,2,3,4,5,6])
b = np.array([0,1,2,3,4,5])

ab_h = np.hstack((a,b))
ab_v = np.vstack((a,b))
ab = np.dot(a,b)
AA, BB = np.meshgrid(a,b)

このように，shape_commentatorをモジュールとして実行します．

python -m shape_commentator src.py

すると，ファイル名.commented.py(今回はsrc.py.commented.py)というファイルができていて，コメントが付いています．やったね．

import numpy as np
a = np.array([1,2,3,4,5,6])  #_ (6,),
b = np.array([0,1,2,3,4,5])  #_ (6,),

ab_h = np.hstack((a,b))  #_ (12,),
ab_v = np.vstack((a,b))  #_ (2, 6),
ab = np.dot(a,b)  #_ (),
AA, BB = np.meshgrid(a,b)  #_ ((6, 6),(6, 6),)

PyTorchで使ってみる

このツールの強力さをもっと知ってもらうために，PyTorchでの例も書きます． PyTorchのコードはkotiraのPyTorch Tutorialから引用します． github.com

ここでクイズです．下のコードの[HERE]の部分は最後の全結合層(FC層)の入力ノード数を表ていますが，これは，FC層に入るテンソルの次元数と合わせなければいけません．さて，この[HERE]の部分には何の数字が入るでしょうか?

PyTorchに触れたことがない方向けに説明すると，layer2の出力は(バッチ(一度に入力する画像のまとまり)サイズ, ノード数(=32), 特徴マップ(中間で出力される画像)の縦の大きさ, 特徴マップの横の大きさ)で，ConvNetに入力画像のバッチを入れた際にforward関数が画像のバッチを引数として呼び出されます．
forward関数の中では，layer2の出力を(バッチサイズ, 残りの次元数を全部まとめた数)のShapeに変えています．fc層の引数のShapeはどうなっているでしょう?

class ConvNet(nn.Module):
    def __init__(self, num_classes=10):
        super(ConvNet, self).__init__()
        self.layer1 = nn.Sequential(
            nn.Conv2d(1, 16, kernel_size=5, stride=1, padding=2),
            nn.BatchNorm2d(16),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2))
        self.layer2 = nn.Sequential(
            nn.Conv2d(16, 32, kernel_size=5, stride=1, padding=2),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2))
        self.fc = nn.Linear(      [    HERE    ]    , num_classes)
        
    def forward(self, x):
        out = self.layer1(x)
        out = self.layer2(out)
        out = out.reshape(out.size(0), -1)
        out = self.fc(out)
        return out

さて，shape_commentatorを利用して全ての代入文に対してshapeを調べましょう．とりあえずnn.Linearに渡す引数の[HERE]は適当な数字で埋めておいて，今までと同様にマジックコマンドを実行します．

さて，どうなるでしょう?

In [2]: %%shape
   ...: import torch 
   ...: import torch.nn as nn
(データローダー部分は省略．)
   ...: class ConvNet(nn.Module): 
   ...:     def __init__(self, num_classes=10): 
   ...:         super(ConvNet, self).__init__() 
   ...:         self.layer1 = nn.Sequential( 
   ...:             nn.Conv2d(1, 16, kernel_size=5, stride=1, padding=2), 
   ...:             nn.BatchNorm2d(16), 
   ...:             nn.ReLU(), 
   ...:             nn.MaxPool2d(kernel_size=2, stride=2)) 
   ...:         self.layer2 = nn.Sequential( 
   ...:             nn.Conv2d(16, 32, kernel_size=5, stride=1, padding=2), 
   ...:             nn.BatchNorm2d(32), 
   ...:             nn.ReLU(), 
   ...:             nn.MaxPool2d(kernel_size=2, stride=2)) 
   ...:         self.fc = nn.Linear( 1,  num_classes) 
   ...:          
   ...:     def forward(self, x): 
   ...:         out = self.layer1(x) 
   ...:         out = self.layer2(out) 
   ...:         out = out.reshape(out.size(0), -1) 
   ...:         out = self.fc(out) 
   ...:         return out 
   ...:  
   ...: model = ConvNet(num_classes).to(device) 
   ...:  
   ...: # Loss and optimizer 
   ...: criterion = nn.CrossEntropyLoss() 
   ...: optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate) 
   ...:  
   ...: # Train the model 
   ...: total_step = len(train_loader) 
   ...: for epoch in range(num_epochs): 
   ...:     for i, (images, labels) in enumerate(train_loader): 
   ...:         images = images.to(device) 
   ...:         labels = labels.to(device) 
   ...:          
   ...:         # Forward pass 
   ...:         outputs = model(images) 
   ...:         loss = criterion(outputs, labels) 
   ...:          
   ...:         # Backward and optimize 
   ...:         optimizer.zero_grad() 
   ...:         loss.backward() 
   ...:         optimizer.step() 
   ...:          
   ...:         if (i+1) % 100 == 0: 
   ...:             print ('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'  
   ...:                    .format(epoch+1, num_epochs, i+1, total_step, loss.item())) 
   ...:                                                                                                                                                           
---------------------------------------------------------------------------
(データローダー部分は省略．以下，出力されたコメント付きスクリプト)                                                                                                                         
# Convolutional neural network (two convolutional layers)
class ConvNet(nn.Module):
    def __init__(self, num_classes=10):
        super(ConvNet, self).__init__()
        self.layer1 = nn.Sequential(  #_ <class 'torch.nn.modules.container.Sequential'>,
            nn.Conv2d(1, 16, kernel_size=5, stride=1, padding=2),
            nn.BatchNorm2d(16),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2))
        self.layer2 = nn.Sequential(  #_ <class 'torch.nn.modules.container.Sequential'>,
            nn.Conv2d(16, 32, kernel_size=5, stride=1, padding=2),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2))
        self.fc = nn.Linear(  1 , num_classes)  #_ <class 'torch.nn.modules.linear.Linear'>,
        
    def forward(self, x):
        out = self.layer1(x)  #_ torch.Size([100, 16, 14, 14]),
        out = self.layer2(out)  #_ torch.Size([100, 32, 7, 7]),
        out = out.reshape(out.size(0), -1)  #_ torch.Size([100, 1568]),
        out = self.fc(out)
        return out

model = ConvNet(num_classes).to(device)  #_ <class '__main__.ConvNet'>,

# Loss and optimizer
criterion = nn.CrossEntropyLoss()  #_ <class 'torch.nn.modules.loss.CrossEntropyLoss'>,
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)  #_ <class 'torch.optim.adam.Adam'>,

# Train the model
total_step = len(train_loader)  #_ <class 'int'>,
for epoch in range(num_epochs):
    for i, (images, labels) in enumerate(train_loader):
        images = images.to(device)  #_ torch.Size([100, 1, 28, 28]),
        labels = labels.to(device)  #_ torch.Size([100]),
        
        # Forward pass
        outputs = model(images)
        loss = criterion(outputs, labels)
        
        # Backward and optimize
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        if (i+1) % 100 == 0:
            print ('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}' 
                   .format(epoch+1, num_epochs, i+1, total_step, loss.item()))
---------------------------------------------------------------------------
(スタックトレースは省略．以下，エラー文)
RuntimeError: size mismatch, m1: [100 x 1568], m2: [1 x 10] at /Users/soumith/code/builder/wheel/pytorch-src/aten/src/TH/generic/THTensorMath.cpp:2070

変数outのshapeは，torch.Size([100, 1568])だと出てきましたね．答えは1568(=32*32*7)です．これで書き換えれば，何も面倒な計算をしたりprint文を入れていちいち調べなくても，ニューラルネットワークの配列のshapeを揃えることができるわけです．どうでしょう，shape_commentatorの便利さを分かっていただけたでしょうか?

その他の機能

In [1]: import shape_commentator                                                                                                     

In [2]: %%shape 
   ...: a = "aaa,bbb"  
   ...: b = a.split(",") 
   ...: c = len(b[2]) 
   ...: d = b[0] 
   ...:  
   ...:                                                                                                                              
a = "aaa,bbb"   #_ str
b = a.split(",")  #_ [str,str,]
c = len(b[2])
d = b[0]
---------------------------------------------------------------------------
(スタックトレースは省略．)
IndexError: list index out of range

途中でエラーになるコードでもOK

commentの実行時にエラーで落ちても，途中までコメントが書かれたソースコードが出力されます．
エラーの原因が分かりやすくなるかと思います．

NumPy以外でも，一般のPythonコードにも利用できる

代入されたオブジェクトに.shapeというattributeがなければ，クラス名をコメントとして出力します．ただし，他のプログラミングなどに使うには，listが展開されてしまうので見づらくなる場面もあるかもしれません．

引数にも対応

引数のあるプログラムでも，

python -m shape_commentator src.py arg1 arg2

と言う様にファイル名の後に引数を入れていけば同じように使えます．

コメントの上書き

shape_commentatorを使ってコメントを付けたソースコードを変更して，shape_commentatorを実行すると，コメントの部分が正しいShapeに書き換わったものが出力されます．これで，ソースコードにShapeを書きっぱなしにしてももう一度shape_commentatorをかければコメントが正しいものになるわけです．
言い換えると，コメントに書いてあるShapeが嘘になってしまうことを防げるわけです．

気づいた方もいると思いますが，Shapeのコメントが通常のPythonのコメントの#とは違って#_になっています．
これは単に実装の都合で，コメントの上書きをしようとするときに#のコメントを使っていると通常のコメントまで消してしまいます．(僕はshape_commentatorの開発途中に，頑張って書いていたコメントを消してしまいました．)
それを防ぐために，#_という#の後に余計な文字を付けてコメントをしていますが，これは今後のリリースで変更するかもしれません．

開発に関して

今後の方針

~~Python2系対応~~ 対応済み
- レガシーなコードを読み解くのに有用なツールにしたいため
PEP 0484の型コメントとの互換対応
- できるのか・・・?
~~TestPyPI，PyPIへのアップロードのためのCI導入~~ 対応済み
- Python2系のテストも回す
~~サイズが大きなlist，tupleでコメントが長くなりすぎないようにする~~ 対応済み
TensorFlowなど他のライブラリでの検証

実装 (ありがとうセキュリティキャンプ)

shape_commentatorの実装に関しては踏み込めませんでしたが，そのうち書きたいと思います．ざっくり言うと，ASTを書き換えてshapeを叩いて情報を取得し，それを使ってソースコードにコメントを書き足しています．簡単にASTを書き換えることができるというのは，セキュリティキャンプ2018 の講義で教わった内容で，講義を受けながらワクワクしてshape_commentatorの構想に思いを巡らせていました．素晴らしい機会をくださった運営の方，講師の方には本当に感謝しています．

ご意見など

shape_commentatorに関する意見，要望などは，Twitter: @_6v_ にリプやDMをくださればすぐに見ます．ドキュメントが足りていないので，TensorFlowで使えた，Pythonのどのバージョンで動かない，お前のやり方は間違っているなど，何でもご意見お待ちしています．

感想はこのブログやはてなブックマークにコメントをくださると嬉しいです．

GitHubにコードを上げてあるので，気に入った方は⭐️付けてもらえると嬉しいです． github.com

CAMPHOR- Advent Calendar 2018 の明日の担当は tori です．おたのしみに〜！