- pytorchのadd_moduleを使用し、任意の場所に層を追加しいたいです。

Question

実現したいこと

• pytorchのadd_moduleを使用し、画像の①の場所に層を追加しいたいです。

発生している問題・エラーメッセージ

add_moduleを使用すると②の場所に追加されてしまいます。ネットで調べたのですが、モデルの最初に追加しているものなどはなく、実装できずに困っています。どのように行えばよいでしょうか。

Comments

[ps_aux_grep]

必ずadd_moduleを使うという制約があるのでしょうか？
また，図示する際は画像で構いませんがコード本体は編集機能上部メニューにあるコードブロック</>の記号から出力される3点バッククォート内部に記述願います．

Answer 1

add_moduleじゃなくても良い場合，というかadd_moduleは後方にしか結合できないので次の選択肢をとりましょう．

from torch import nn

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(1000, 100)
        self.fc2 = nn.Linear(100, 10)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(0.2)
    
    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.dropout(x)
        x = self.fc2(x)
        return x

class IOWrappedNet(nn.Module):
    def __init__(self, net):
        super().__init__()
        self.fc0 = nn.Linear(10000, 1000)
        self.net = net
        self.softmax = nn.Softmax()
    
    def forward(self, x):
        x = self.fc0(x)
        x = self.net(x)
        x = self.softmax(x)
        return x

net = IOWrappedNet(Net())
print(net)

IOWrappedNet(
  (fc0): Linear(in_features=10000, out_features=1000, bias=True)
  (net): Net(
    (fc1): Linear(in_features=1000, out_features=100, bias=True)
    (fc2): Linear(in_features=100, out_features=10, bias=True)
    (relu): ReLU()
    (dropout): Dropout(p=0.2, inplace=False)
  )
  (softmax): Softmax(dim=None)
)