配列の形状(Shape)のコメントを付けるJupyter Notebook拡張を作った

概要

下の図のように,Pythonのコードに配列の形状(Shape)だったりクラスだったりをコードに付け加えるJupyter Notebook拡張を作りました. 今度こそ超便利.

f:id:shiba6v:20190804173242p:plain

GitHubリポジトリはこちらです. github.com

はじめに

以前,NumPyやPyTorchで使える超便利ツールを作った の記事でShape Commentatorを作ったことを紹介しました.

このShape Commentatorは実行時の情報を抜き出すという都合上,使い方が煩雑でした.(ブログの説明がかなり長いのはそういうことです・・・)

このせいで自分でも本当に困った時しか使わなくなってしまい,もっと便利にしたい,できればJupyter Notebookの拡張機能として使えるようにしたいと思って作りました.

インストール

Shape Commentatorをpipでインストールし,Jupyter Notebook 拡張としてインストールします.

nbextensionsが入っていない場合は,jupyter_contrib_nbextensions 公式ドキュメントからインストールしてください.

pip install -U shape_commentator
jupyter nbextension install https://github.com/shiba6v/jupyter-shape-commentator/archive/master.zip --user

そして,nbextensionsの設定からJupyter Shape CommentatorをOnにします. nbextensions

これで準備は完了です.

使う

コメントをつけたいセルを選んで, Shape のボタンを押すだけでOKです!

コメントを消したいときは Shape のボタンを押せばOKです.消えるコメントはShape Commentatorで付けたコメントだけで,通常のコメントが消えることはありません.

f:id:shiba6v:20190804173242p:plain

補足

Shape Commentatorで付けたコメントは,通常のコメントとの区別のために「 #_ 」 というように「 # 」の後に「 _ 」を入れています. いずれはPEP484に準拠したいですが,一旦このように妥協しています.

デモ

例えば,ニューラルネットワークのコードでforwardの中のShapeがわからなくなっても簡単に見ることができます.

デモはPyTorchで書いていますがPandasやChainerでも使えます."shape"というattributeが生えていれば何でも良いのです.

今後

今回Jupyter Notebook拡張を書いたことでやっと楽に使えるレベルまで達したと思います.(個人の感想です.)

今後も機能を追加していけたらいいと思っています.

Shape Commentator使った!などの感想をTwitterしてくださった方,バグ報告のissueが立ててくださった方,感想をくれたお友達に感謝します.

日本語のJupyter NotebookをPDFとしてダウンロードする裏技

Jupyter NotebookをPDF化する際にLaTeXを入れて頑張る方法がありますが,もっと楽をする方法があります.

ブラウザの機能を使って印刷すれば良いんです.

f:id:shiba6v:20190627235218p:plain
notebook

Ctrl+P (Macの場合はCmd+P) で印刷ダイアログを出して,「PDFに保存」すれば完成です. 簡単ですね.

f:id:shiba6v:20190627235234p:plain
pdf

「Jupyter Notebook PDF」で調べてもあまりこのやり方が出てこなかったので,ブログとして書いてみました.

結構このやり方を知らない人が多いので,この記事が検索で上位に来れば救われる人が増えるのかなと思います・・・

toxからbatsを使ってPythonのバージョンを変えながらシェルスクリプトでテストする際のエラー

こんにちは. この前リリースしたPythonで使えるツールのshape_commentatorの評判が良く,最近はCIでPyPIにパッケージをアップロドする設定を書いたりしてCIとわいわい遊んでいます.

shiba6v.hatenablog.com

tox と bats

テストには,doctestに加えてbats を利用しています.これは,シェルスクリプト でテストを書けるツールで,今回のようなファイルを直接出力するなどPythonのコードよりもシェルスクリプトでテストを書きたいようなもののテストに適していると思い,導入しました.

shape_commentatorは,過去の数値計算のコードを読み解くツールとして有用になる可能性があるため,様々なPythonのバージョンでテストしておきたいという気持ちがあります.Pythonの様々なバージョンでのテストと言えばtoxですが,これはvirtualenvを使って様々なバージョンのPythonに対してテストを行うツールです.

tox.iniの書き方

tox.iniからbatsを呼ぶために,このようにtox.iniを書きました.

[tox]
envlist = py36
[testenv]
deps = numpy
whitelist_externals = bats
commands = bats --tap tests/test_all.bats

気をつける点

注意点は以下の2点です. 1. toxからbatsのような外部コマンドを呼ぶには,whitelist_externalsに書く必要があります. 2. batsを直接ターミナルから呼び出さない場合にあたるので,--tapオプションを付けます.

1をやらないと,

WARNING: test command found but not installed in testenv
  cmd: /usr/local/bin/bats
  env: /Users/xxxxxx/shape_commentator/.tox/py36
Maybe you forgot to specify a dependency? See also the whitelist_externals envconfig setting.

といったエラーが出ます. toxのcommandsからpython以外のコマンド(bats)を呼び出しているために出るエラーです.

2をやらないと,

tput: No value for $TERM and no -T specified
ERROR: InvocationError for command '/usr/local/bin/bats tests/test_all.bats' (exited with code 2)

といったエラーが出ます. toxからbatsを実行すると,batsが利用しているtputコマンドから出るエラーです.

まとめると,tox.iniでbatsを実行する際は,whitelist_externals--tapを書きましょうということです.

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 です. おたのしみに〜!

WSLでCuda toolkitがインストールできない問題について

WSLでCuda toolkitを入れようとして,このページを見ながらやっていたのですが,

qiita.com

sudo apt-key adv --fetch-keys http://developer.download.nvidia.com/compute/cuda/repos/ubuntu1604/x86_64/7fa2af80.pub
の部分で gpg: connecting dirmngr failed: IPC connect call failed
といったエラーが出ました.

エラー文でググると,WSLのバグのようで,
curl -sL http://developer.download.nvidia.com/compute/cuda/repos/ubuntu1604/x86_64/7fa2af80.pub | sudo apt-key add
で行けました.

よかった~

参考
https://github.com/Microsoft/WSL/issues/3286

PyTorchのDataParallelのモデルを保存する

PyTorchで複数GPUで学習させる場合,

model = nn.DataParallel(model, device_ids=[0,1,2])

のようにDataParallelで保存しますが,このモデルを保存したい場合にcuda runtime error : out of memoryが出ることがあります.
その場合は,下のようにDataParallelから元のモデルを取り出してCPUのモデルに変えてあげることで保存できるようになります.

torch.save(model.module.cpu(),file_path)

読み込み時はこうすればOK

new_model = torch.load(file_path)

参考

Optional: Data Parallelism — PyTorch Tutorials 1.0.0.dev20181002 documentation

cProfileを使ってPythonでの計算時間を測定する

最近は,4回生になり研究室に配属されて画像処理の勉強をしています. 論文を読んでそれを実装してみるという流れで勉強していますが,REPLでnumpyの行列計算のコードを打っていると,目立って時間がかかる行があることが分かります. そこで,繰り返し呼ぶ処理や重い処理にかかる時間を,定量的に調べたいと思いました.

やりたいこと

  • スクリプト全体で,時間がかかっている箇所を調べたい
  • 遅い処理を見つけて高速化したい
  • できるだけ簡単にプロファイリングしたい

環境・使用ライブラリ

  • Ubuntu 16.04.4 LTS
    • WSL上で動かしています
  • Python 3.5.2
  • cPofile
  • pstats

Pythonで使えるプロファイラについて調べてみると,cProfileとline_profilerという2つのプロファイラがあることが分かりました. line_profilerを使う場合は関数に@profileというデコレータを追加する必要があり,コードを書き換えずにREPLで試したいので今回はパスしました.Pythonに標準で入っているcProfileを使うことにしました.pstatsはプロファイリング結果を表示できるモジュールです.

方法

時間を計測したいコードの前に以下のコードを貼り付けて,プロファイラを有効にします.

import cProfile
import pstats
c_profile = cProfile.Profile()
c_profile.enable()

計測したいコードの後に以下のコードを貼り付けます. print_stats(10)はtottime(そのメソッドにかかった時間の合計)の上位10個をとってくるという意味です.

c_profile.disable()
c_stats = pstats.Stats(c_profile)
c_stats.sort_stats('tottime').print_stats(10)

全体をプロファイラにかける

結果

書いていたスクリプト(本来は貼るべきですが,説明すると長くなりそうなので省略)をプロファイラにかけると,6.8秒もかかっていて,そのうちのほとんどがnumpy/linalg/linalg.py:1299(svd)にかかっていることが分かりました.

        267323 function calls (261982 primitive calls) in 6.844 seconds

   Ordered by: internal time
   List reduced from 2743 to 10 due to restriction <10>

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
        1    4.634    4.634    4.634    4.634 /usr/local/lib/python3.5/dist-packages/numpy/linalg/linalg.py:1299(svd)
        1    1.082    1.082    1.082    1.082 {built-in method numpy.core.multiarray.matmul}
       16    0.195    0.012    0.195    0.012 {imread}
       28    0.116    0.004    0.121    0.004 {built-in method _imp.create_dynamic}
      286    0.083    0.000    0.099    0.000 <frozen importlib._bootstrap_external>:816(get_data)
        1    0.078    0.078    5.793    5.793 /usr/local/lib/python3.5/dist-packages/numpy/linalg/linalg.py:1652(pinv)
      286    0.062    0.000    0.062    0.000 {built-in method marshal.loads}
1273/1218    0.027    0.000    0.119    0.000 {built-in method builtins.__build_class__}
     1269    0.022    0.000    0.022    0.000 {built-in method posix.stat}
      286    0.016    0.000    0.016    0.000 {method 'read' of '_io.FileIO' objects}

考察

コード中に x = np.dot( np.linalg.pinv(A) , b) という部分があり,そこがnp.linalgを使っていたので遅そうだと考えました. この部分はこの記事の本題とは関係ありませんが,  A {\bf x} = {\bf b} を解いています. 優決定系(変数の数より連立方程式の数が多い)の行列で最小二乗誤差が小さくなるようなxを求める式です.詳しくはこちらのサイトに書いてあります.

最小二乗法とフィッティングとモデルパラメータ推定(アマチュア用)

行列計算を改善する

先ほど遅かった部分の計算時間を見るために,次のコードをREPLに貼り付けました.

import cProfile
import pstats
c_profile = cProfile.Profile()
c_profile.enable()

x = np.dot( np.linalg.pinv(A) , b)

c_profile.disable()
c_stats = pstats.Stats(c_profile)
c_stats.sort_stats('tottime').print_stats(10)

結果

         51 function calls in 5.461 seconds

   Ordered by: internal time
   List reduced from 29 to 10 due to restriction <10>

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
        1    4.353    4.353    4.353    4.353 /usr/local/lib/python3.5/dist-packages/numpy/linalg/linalg.py:1299(svd)
        1    1.026    1.026    1.026    1.026 {built-in method numpy.core.multiarray.matmul}
        1    0.079    0.079    5.458    5.458 /usr/local/lib/python3.5/dist-packages/numpy/linalg/linalg.py:1652(pinv)
        1    0.003    0.003    0.003    0.003 {built-in method numpy.core.multiarray.dot}
        1    0.000    0.000    0.000    0.000 {method 'reduce' of 'numpy.ufunc' objects}
        3    0.000    0.000    0.000    0.000 /usr/local/lib/python3.5/dist-packages/numpy/core/numeric.py:424(asarray)
        1    0.000    0.000    0.000    0.000 /usr/local/lib/python3.5/dist-packages/numpy/core/fromnumeric.py:2222(amax)
        3    0.000    0.000    0.000    0.000 {built-in method numpy.core.multiarray.array}
        3    0.000    0.000    0.000    0.000 {method 'astype' of 'numpy.ndarray' objects}
        1    0.000    0.000    0.000    0.000 <stdin>:1(<module>)

考察

ここでやはり4秒程度かかっていたことが分かります. これまでは, A {\bf x} = {\bf b} を解くときに[tex: AT A {\bf x} = AT {\bf b}] から疑似逆行列(np.linalg.pinv(A))を求めて[tex: {\bf x} = (AT A)^{-1} AT {\bf b}] を計算して解いていました. それをやめて,lstsqという最小二乗法のメソッドを使ってnp.linalg.lstsq(A.T @ A, A.T @ b)とする方が効率的かもしれないと思いました.

改善して比較する

先ほどと比較するために,次のコードをREPLに貼り付けました.

import cProfile
import pstats
c_profile = cProfile.Profile()
c_profile.enable()

AT =A.T
x,res,rank,s = np.linalg.lstsq(AT @ A, AT @ b)

c_profile.disable()
c_stats = pstats.Stats(c_profile)
c_stats.sort_stats('tottime').print_stats(10)

結果

         122 function calls in 0.389 seconds

   Ordered by: internal time
   List reduced from 46 to 10 due to restriction <10>

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
        2    0.383    0.192    0.383    0.192 {built-in method numpy.linalg.lapack_lite.dgelsd}
        2    0.003    0.002    0.003    0.002 {built-in method numpy.core.multiarray._fastCopyAndTranspose}
        1    0.001    0.001    0.001    0.001 {method 'write' of '_io.TextIOWrapper' objects}
        8    0.000    0.000    0.000    0.000 {built-in method posix.stat}
        1    0.000    0.000    0.389    0.389 /usr/local/lib/python3.5/dist-packages/numpy/linalg/linalg.py:1880(lstsq)
        5    0.000    0.000    0.000    0.000 {built-in method numpy.core.multiarray.zeros}
        1    0.000    0.000    0.000    0.000 /usr/lib/python3.5/linecache.py:82(updatecache)
        1    0.000    0.000    0.002    0.002 {built-in method _warnings.warn}
        7    0.000    0.000    0.000    0.000 /usr/lib/python3.5/posixpath.py:71(join)
        1    0.000    0.000    0.000    0.000 {method 'reduce' of 'numpy.ufunc' objects}

は,速い・・・

考察

4秒以上かかっていた先ほどと比べて,0.389秒とかなり速くなっていることが分かりました.

感想

REPLでもプロファイラを使って実行時間を比較できることが分かりました. 処理が遅いときに原因となっている箇所を探すのに使えそうです. CliborやClipyのスニペットに追加しておくと便利そうだと思いました. メソッドごとに見られるline_profilerも便利そうなので,cProfileに飽きてきたらそちらも使ってみようと思います.