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])