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

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

はじめに

こんにちは.シバニャンです.最近は卒論で画像系の研究をしていますが,機械学習の理論的なところを突き詰めずにやってきたツケが回ってきて苦労しています😢

卒業研究では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_commentator.comment(In[len(In)-2], globals(), locals()) 

という一行のプログラムを実行します.

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

In [1]: 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)                                                                                                                                                                                                     

In [2]: import shape_commentator 
   ...: shape_commentator.comment(In[len(In)-2], globals(), locals()) 
   ...:                                                                                                                                      
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のコメントが書いてあったら見る側も嬉しいと思います.

In[2]の二行目を解説すると,これはshape_commentatorのcommentというメソッドを 呼び出していて,引数1にはソースコード,引数2にはグローバル変数,引数3にはローカル変数を入れます. InはIPythonの入力で,In[len(In)-2]は一つ前の入力(ここではIn[1])を表しています. つまり,In[1]のソースコードと現在の変数をcommentメソッドに渡しているということです.

これによってshape_commentatorは,渡されたソースコード実際に実行し代入時にshapeの情報を取り出してソースコードに付け加えて出力します. In[1]は途中で実行を中断してもIn[1]にソースコードが入るのでOKです.

In[1]の最後の行の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のコードはyunjeyさんの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を調べましょう. とりあえず[HERE]は適当な数字で埋めておいて,Shapeを知りたい部分のコードを実行した後(途中で止めてOK),今までと同様にコメントするコードを実行します.

import shape_commentator
shape_commentator.comment(In[len(In)-2], globals(), locals())

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

In [2]: # Convolutional neural network (two convolutional layers) 
   ...: 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())) 
   ...:                                                                                                                                                           
---------------------------------------------------------------------------
(省略)
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

In [3]: import shape_commentator 
   ...: shape_commentator.comment(In[len(In)-2], globals(), locals()) 
   ...:                                                                                                                                                           
# 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

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

その他の機能

In [1]: a = "aaa,bbb" 
   ...: b = a.split(",") 
   ...: c = len(b[2]) 
   ...: d = b[0]                                                                                                                             
---------------------------------------------------------------------------
(省略)
IndexError: list index out of range

In [2]: import shape_commentator 
   ...: shape_commentator.comment(In[len(In)-2], globals(), locals())                                                                        
a = "aaa,bbb"  #_ <class 'str'>,
b = a.split(",")  #_ (<class 'str'>,<class '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 です. おたのしみに〜!