forward_one_layer()を実装しよう

CNNのレイヤについて,インデックスで指定してそこまで学習を進める,といった処理を繰り返せば良い?

code: one_layer.py

def forward_one_layer(self, x, layer_idx):

"""

指定した層番号(layer_idx)のConv2d層からPooling層までをまとめてforwardし、出力を返す

"""

# Conv2d層のインデックスを探す

conv_count = -1

start_idx = None

end_idx = None

for idx, layer in enumerate(self.cnn):

if isinstance(layer, nn.Conv2d):

conv_count += 1

if conv_count == layer_idx:

start_idx = idx

if start_idx is not None and isinstance(layer, nn.AvgPool2d):

end_idx = idx

break

if start_idx is None or end_idx is None:

raise ValueError(f"Could not find Conv2d and Pooling for layer {layer_idx}")

out = x

for idx in range(start_idx, end_idx + 1):

return out