【扒代码】代码调试TransformerEncoder

    def forward(self, src, pos_emb, src_mask, src_key_padding_mask):
        # src = torch.Size([4096, 4, 256])
        # pos_emb = torch.Size([4096, 4, 256])
        output = src
        for layer in self.layers:
            print("layer : ",layer)
            output = layer(output, pos_emb, src_mask, src_key_padding_mask)
            print("output : ",output.shape)
        return self.norm(output)

layer :  TransformerEncoderLayer(
  (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
  (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
  (dropout1): Dropout(p=0.1, inplace=False)
  (dropout2): Dropout(p=0.1, inplace=False)
  (self_attn): MultiheadAttention(
    (out_proj): NonDynamicallyQuantizableLinear(in_features=256, out_features=256, bias=True)
  )
  (mlp): MLP(
    (linear1): Linear(in_features=256, out_features=2048, bias=True)
    (linear2): Linear(in_features=2048, out_features=256, bias=True)
    (dropout): Dropout(p=0.1, inplace=False)
    (activation): GELU(approximate='none')
  )
)
output :  torch.Size([4096, 4, 256])
layer :  TransformerEncoderLayer(
  (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
  (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
  (dropout1): Dropout(p=0.1, inplace=False)
  (dropout2): Dropout(p=0.1, inplace=False)
  (self_attn): MultiheadAttention(
    (out_proj): NonDynamicallyQuantizableLinear(in_features=256, out_features=256, bias=True)
  )
  (mlp): MLP(
    (linear1): Linear(in_features=256, out_features=2048, bias=True)
    (linear2): Linear(in_features=2048, out_features=256, bias=True)
    (dropout): Dropout(p=0.1, inplace=False)
    (activation): GELU(approximate='none')
  )
)
output :  torch.Size([4096, 4, 256])
layer :  TransformerEncoderLayer(
  (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
  (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
  (dropout1): Dropout(p=0.1, inplace=False)
  (dropout2): Dropout(p=0.1, inplace=False)
  (self_attn): MultiheadAttention(
    (out_proj): NonDynamicallyQuantizableLinear(in_features=256, out_features=256, bias=True)
  )
  (mlp): MLP(
    (linear1): Linear(in_features=256, out_features=2048, bias=True)
    (linear2): Linear(in_features=2048, out_features=256, bias=True)
    (dropout): Dropout(p=0.1, inplace=False)
    (activation): GELU(approximate='none')
  )
)
output :  torch.Size([4096, 4, 256])