PyTorchを使ったネットワーク（20行程度）の解釈を手伝ってほしいです

Question

助けて…

PyTorchを使ったCNNの勉強をしています．
このAttentionに関する記事の，Pytorchを使ったネットワークの定義が理解できません．
助けてほしいです…．

(これまで既存のコードのコピペしかしていなかったので，ネットワークの次元の変化が追えていないです????)

理解できないところ

forward関数です．

ResNet34の出力層を取ったものを使っているので，forward関数にてx = self.features(x)で512×3x3の特徴量xを得ていると思います．
次にself.attn_conv(x)でnn.Sequential(nn.Conv2d(512,1,1), nn.Sigmoid())に通用していると思います．
この出力が（コメントより），[B, 1, H, W]次元となるのが理解できていないところです．

512×3×3 => B×1×H×W
となるのが追えません．どなたかご解説願えないでしょうか…!!

該当のネットワーク，ソースコード

class SimpleAttentionNetwork(nn.Module):
        def __init__(self, num_classes):
            super().__init__()

            base_model = resnet34(pretrained=True)
            self.features = nn.Sequential(*[layer for layer in base_model.children()][:-2])
            self.attn_conv = nn.Sequential(
                nn.Conv2d(512, 1, 1),
                nn.Sigmoid()
		    )
            self.fc = nn.Sequential(
                nn.Dropout(0.5),
                nn.Linear(512, num_classes)
            )
            self.mask_ = None

        def forward(self, x):
            x = self.features(x)

            attn = self.attn_conv(x)  # [B, 1, H, W]
            B, _, H, W = attn.shape
            self.mask_ = attn.detach().cpu()

            x = x * attn
            x = F.adaptive_avg_pool2d(x, (1, 1))
            x = x.reshape(B, -1)

            return self.fc(x)

参考

深層学習入門：画像分類（5）Attention機構｜SBテクノロジー（SBT)
ResNet（Residual Network)の実装｜AIdrops

Comments

[jbpb0]

> これまで既存のコードのコピペしかしていなかったので，ネットワークの次元の変化が追えていないです

その自覚があるのでしたら、まずそこからやりましょうよ

こういうツールもあるし
https://qiita.com/tand826/items/ae3349495944048fd120

[r.kanke]

やはりそうですか…．正直どう手をつけたものかと敬遠していました．
向き合う時が来たんですね…．

torchinfo入れてみました．こんなに見やすくなるんですね！
ありがとうございます！

[quickquip]

BやHやWの実際の値を確認したけどどうしてその値になるのかわからない、という質問なのか
BやHやWが実際にどんな値になるのかわからない、という質問なのか
どうしてBやHやWに代入するのかわからない、という質問なのか
BやHやWってなんですか、という質問なのか
の区別が付かないですね……

[r.kanke]

Bが何を表すのかが知りたかったです．

しかし指摘して頂いた内，実際の値は確認していませんでした．やってきます．

[quickquip]

だとすると質問は「Conv2Dの中身を理解したい」と同義ですね。
Conv2Dで変換した結果の、テンソルの第1要素(インデックス0番）の次元数をBとしよう
というだけなので。 追記:入れ違いました
追記: 次元数 は誤用でした

[r.kanke]

なるほどそこの理解が足りなかったのですね…．なんで詰まったのかが分かったので有り難いです！
Conv2dの復習をしようと思います．

Answer 1

Attentionマスクのサイズ（次元）を勘違いしていました．

画像では入力画像と同じサイズになっていますが，それは可視化した場合．実際にはバッチサイズ×チャネル×タテ×ヨコ = 32×1×7×7となっていました．

xとattnのshapeを画面出力させて確認しました．　　
x.shape: torch.Size([32, 512, 7, 7])
attn.shape: torch.Size([32, 1, 7, 7])

Comments

[r.kanke]

@quickquipさんがご指摘くださったように，何が分からないのか，明確にすると対処法が見つかりそうです．
@jbpb0さん，@quickquipさん，アドバイスありがとうございました．非常に助かりました????????