PyTorch/libtorch の load/save メモ(主にデータコンバート向け)

## 背景

* 学習は GPU, 推論は CPU で行うなどで, torch.load や torch.save の仕組みを知りたい
* weight のデータをちょこっと編集したりとか, weight を自前 C++ 機械学習アプリなどで読み込みたい

## ソースコードツアー

https://github.com/pytorch/pytorch

python レイヤー(`torch.load`, `torch.save`)は, `torch/serialization.py` に実装があります.

## libtorch(C++)関連

`torch/csrc` にあります.

https://github.com/pytorch/pytorch/tree/master/torch/csrc

##  シリアライズ

https://github.com/pytorch/pytorch/blob/master/torch/serialization.py

`torch.save`, `torch.load` では, python pickle で Python object(model + weights)をシリアライズします.
これにより, モデルの現状(~= checkpoint. model + weight)を保存, ロードすることができるようになっています.

対応する Python のスクリプトも保存されるので, `torch.load` でモデルを復元した場合は, 元の model のソースコードは不要になります(たぶん)

model の Python object instance(?) を取得する場合は, 通常 `torch.load("checkpoint.pt")['model']` で取得できます.

### libtorch 

libtorch(C++)では, Python のサブセットライクな TorchScript 形式でモデルを読めますが, どうも Python コード or Python に近い言語を一から実装して, 処理しているようですね. なかなかの力技です...

weights だけの場合は, NPZ 形式(npy の複数ファイルを zip uncompressed でひとつにまとめたもの)での対応になっています.
 
## map_location

CUDA で学習してシリアライズしたデータだと, CUDA の情報もふくまれているため, それを CPU only の環境で `torch.load` するとエラーになります. これを解決するものとして `map_location` オプション引数があります.
map_location に指定できるものとしてはいろいろパターンがあり, 関数(callable)を指定も可能になっています. 

データコンバートメインであれば, `"cpu"` 文字列, もしくは `torch.device('cpu')` 指定で事足りるかと思います.
 


## state_dict

https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict

モデルの weight データ(parameter)が保存されている Python dictionary(普通の Python での OrderedDict 形式)です.
 
### Weight(パラメータ)の名前

weight にはなにかしら unique な名前がつきます.

https://discuss.pytorch.org/t/how-to-give-pytorch-layer-a-name/5521

```
self.my_name = nn.Linear
```

とすると, `my_name` で Linear の weight に名前がつきます.
配列(nn.ModuleList)の場合は layers.0.my_name, my_layers.1.my_name, ... のように, `.N.` のルールでつくようです.


したがって, train したモデルと, inference(TorchScript, libtorch)でモデルが異なる場合, うまく名前を見てコンバートする必要がああります.
(たとえば現状(v1.4.0)では TorchScript は `nn.ModuleList` の配列アクセスが扱えないので, 手動で展開とか, 別クラスを作って iterate させるときなど)

### load_state_dict

`load_state_dict()` で weight データを読み込めます.
torch.load や, dict を直接指定(重みデータは `torch.Tensor` がベター)でも OK です.

dict は, ネットワークのすべてのレイヤー名に対応する key が設定されている必要があります.
レイヤーの一部の weight だけ変えたい場合は,

```
conv0.weight = torch.Tensor(...)
conv0.bias = torch.Tensor(...)
```

のようにすることができます. ただこれだと shape のチェックが入らないので注意です.
(なにか他にいい方法がありそうな気はするが)

## 事例

checkpoint データから, weight(state_dict) だけ書き出して, 別のモデルへ読み込む(weight を設定する).

```
device = torch.device('cpu')

model = torch.load("pretrained.pt", map_location=device)['model']

torch.save(model.state_dict(), "weights.pt"))

weights = torch.load("weights.pt", map_location=device)

my_model = MyModel()
my_model.load_state_dict(weights)

```

モデルと weight 全体(pickle シリアリゼーション)も, weight(state_dict) だけもどちらも `torch.load` で読み込めるのでやや混乱しますね.


双方のモデルでテンソルの大きさなどが違うとか, MyModel に対応する weight が無いとかだと, エラーを出してくれます.

 
## libtorch 

Python オブジェクトを読み書きする pickle_load, pickle_save があります

https://github.com/pytorch/pytorch/issues/20356

pickle 対応のために, libtorch(C++) に自前 pickle 実装してます! 
https://github.com/pytorch/pytorch/blob/master/torch/csrc/jit/serialization/pickler.h

テンソルだけロード, 保存は現状できなさそうです(なにかしら Model を作る必要がある)

重みのテンソルデータだけ(基本 float32 型を想定)やりとりしたい場合は, NPY, NPZ で Numpy 形式にして, cnpy を使うのがよいでしょう. https://github.com/rogersce/cnpy 

nanosnap にも cnpy でテンソルデータを読む機能があります.

https://github.com/lighttransport/nanosnap