PyTorchのモデルの書き方の整理

統計・機械学習 python

サマリ

いくつかの書き方があるが，混乱するのは下記の種類が色々Webに見られるからか。

  • ① initにはパラメタありだけ，パラメタが無いものはforwardに。これが一番柔軟性があるように思う。モデルの中間層を細かく取り出したり，活性化関数を変えたりしたいならこれが良いと思う。
  • ②initにReLuなども書いてしまう。print(model)でReLuなども含めて出力してくれる点が良い点。
  • ③ Sequencialを使う。特定の塊ごとに意味を持たせられる。

これらは，それぞれ排他的に使う必要はなくて，適材適所でMixして使う。方針としては，学習後に分解することのない特定の塊があればどんどんSequencialに入れて，initにReLuなどのパラメタが無い層も含めて書いて良いと思う。もし，特定のレイヤを置換したりしたい場合は，それは独立したSequentialか，独立したレイヤにする。

基本事項

  • nn.Moduleを継承して、構成要素をinitに定義して順方向をforwardに記載。callを描かなくても、モデルに入力を渡せば自動的にforwardを実行してくれる。
  • Sequentialを使ったり、Activationを関数に渡してしまったり、色々なパタンがあるので整理しておく。

パタン１(最もシンプルなパタン)

  • 学習パラメタを持つものはinitに、関数はforwardに。
  • 公式チュートリアルはこの書き方。(Poolingもパラメタは無いからforward)
  • print(model)で出力されるのはinitに描かれたもの。
class MLP(nn.Module):
    def __init__(self, input_size, hidden_size, num_classes):
        super(MultiLayerPerceptron, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)        
        self.fc2 = nn.Linear(hidden_size, hidden_size)
        self.fc3 = nn.Linear(hidden_size, num_classes)        
    
    def forward(self, x):
        h = F.relu(self.fc1(x))
        h = F.relu(self.fc2(h))
        out = self.fc3(h)
        return out

パタン2(ReLUなどもinitに集約)

  • パラメタ付きか否かの分離ができないが、print文でReLUなども出力されるのでその点は良い。この場合はReLUは関数(F.relu)ではなくクラス(nn.ReLU)を使う。
class MLP(nn.Module):
    def __init__(self, input_size, hidden_size, num_classes):
        super(MultiLayerPerceptron, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)        
        self.ac1 = nn.ReLU()
        self.fc2 = nn.Linear(hidden_size, hidden_size)
        self.ac2 = nn.ReLU()
        self.fc3 = nn.Linear(hidden_size, num_classes)        
    
    def forward(self, x):
        h = self.ac1(self.fc1(x))
        h = self.ac2(self.fc2(h))
        out = self.fc3(h)
        return out

パタン3(Sequential)

  • KerasのSequentialのような書き方が可能。単純なMLPやCNNならこれで良い。一方で多入力、多出力なネットワークの記述は出来ない。
model = nn.Sequential()
model.add_module('fc1', nn.Linear(input_size, hidden_size))
model.add_module('relu1', nn.ReLU())
model.add_module('fc2', nn.Linear(hidden_size, hidden_size))
model.add_module('relu2', nn.ReLU())
model.add_module('fc3', nn.Linear(hidden_size, n_classes))
  • 先にリストに層をまとめておくことも可能。 layers = [nn.Linear(input_size, hidden_size), nn.ReLU(), nn.Linear(hidden_size, hidden_size), nn.ReLU(), nn.Linear(hiddensize, n_classes)] model = nn.Sequential(*layers)