| 107 | return xyz, output, xyz_inds |
| 108 | |
| 109 | class TransformerEncoderLayer(nn.Module): |
| 110 | |
| 111 | def __init__(self, d_model, nhead=4, dim_feedforward=128, |
| 112 | dropout=0.1, dropout_attn=None, |
| 113 | activation="relu", normalize_before=True, norm_name="ln", |
| 114 | use_ffn=True, |
| 115 | ffn_use_bias=True): |
| 116 | super().__init__() |
| 117 | if dropout_attn is None: |
| 118 | dropout_attn = dropout |
| 119 | self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout_attn) |
| 120 | self.use_ffn = use_ffn |
| 121 | if self.use_ffn: |
| 122 | # Implementation of Feedforward model |
| 123 | self.linear1 = nn.Linear(d_model, dim_feedforward, bias=ffn_use_bias) |
| 124 | self.dropout = nn.Dropout(dropout, inplace=True) |
| 125 | self.linear2 = nn.Linear(dim_feedforward, d_model, bias=ffn_use_bias) |
| 126 | self.norm2 = NORM_DICT[norm_name](d_model) |
| 127 | self.norm2 = NORM_DICT[norm_name](d_model) |
| 128 | self.dropout2 = nn.Dropout(dropout, inplace=True) |
| 129 | |
| 130 | self.norm1 = NORM_DICT[norm_name](d_model) |
| 131 | self.dropout1 = nn.Dropout(dropout, inplace=True) |
| 132 | |
| 133 | self.activation = ACTIVATION_DICT[activation]() |
| 134 | self.normalize_before = normalize_before |
| 135 | self.nhead = nhead |
| 136 | |
| 137 | def with_pos_embed(self, tensor, pos: Optional[Tensor]): |
| 138 | return tensor if pos is None else tensor + pos |
| 139 | |
| 140 | def forward_post(self, |
| 141 | src, |
| 142 | src_mask: Optional[Tensor] = None, |
| 143 | src_key_padding_mask: Optional[Tensor] = None, |
| 144 | pos: Optional[Tensor] = None): |
| 145 | q = k = self.with_pos_embed(src, pos) |
| 146 | value = src |
| 147 | src2 = self.self_attn(q, k, value=value, attn_mask=src_mask, |
| 148 | key_padding_mask=src_key_padding_mask)[0] |
| 149 | src = src + self.dropout1(src2) |
| 150 | if self.use_norm_fn_on_input: |
| 151 | src = self.norm1(src) |
| 152 | if self.use_ffn: |
| 153 | src2 = self.linear2(self.dropout(self.activation(self.linear1(src)))) |
| 154 | src = src + self.dropout2(src2) |
| 155 | src = self.norm2(src) |
| 156 | return src |
| 157 | |
| 158 | def forward_pre(self, src, |
| 159 | src_mask: Optional[Tensor] = None, |
| 160 | src_key_padding_mask: Optional[Tensor] = None, |
| 161 | pos: Optional[Tensor] = None, |
| 162 | return_attn_weights: Optional [Tensor] = False): |
| 163 | |
| 164 | src2 = self.norm1(src) |
| 165 | value = src2 |
| 166 | q = k = self.with_pos_embed(src2, pos) |