[코드 공유] Adaptive Positional 2D Encoding 구현입니다

[shjas94]

class AdaptivePositionalEncoding2D(nn.Module):
    def __init__(self, in_channels, max_h=64, max_w=128, dropout=0.1):
        super(AdaptivePositionalEncoding2D, self).__init__()

        self.h_position_encoder = self.generate_encoder(
            in_channels // 2, max_h)
        self.w_position_encoder = self.generate_encoder(
            in_channels // 2, max_w)

        self.h_linear = nn.Linear(in_channels // 2, in_channels // 2)
        self.w_linear = nn.Linear(in_channels // 2, in_channels // 2)

        self.avgpool = nn.AdaptiveAvgPool2d(1)

        self.inter = nn.Linear(in_channels, in_channels)
        self.relu = nn.ReLU(inplace=True)
        self.alpha = nn.Linear(in_channels, in_channels)
        self.sigmoid = nn.Sigmoid()

        self.dropout = nn.Dropout(p=dropout)

    def generate_encoder(self, in_channels, max_len):
        pos = torch.arange(max_len).float().unsqueeze(1)
        i = torch.arange(in_channels).float().unsqueeze(0)
        angle_rates = 1 / torch.pow(10000, (2 * (i // 2)) / in_channels)
        position_encoder = pos * angle_rates
        position_encoder[:, 0::2] = torch.sin(position_encoder[:, 0::2])
        position_encoder[:, 1::2] = torch.cos(position_encoder[:, 1::2])
        return position_encoder  # (Max_len, In_channel)

    def forward(self, input):
        # Require DEBUG
        b, c, h, w = input.size()
        # adaptive
        e = self.avgpool(input).squeeze()
        a = self.sigmoid(self.alpha(
            self.relu(self.inter(e))))
        a = a.reshape((-1, 2, 1, c//2))

        ##########

        h_pos_encoding = (
            self.h_position_encoder[:h, :].unsqueeze(1).to(input.get_device())
        )

        h_pos_encoding = self.h_linear(h_pos_encoding)  # [H, 1, D]

        w_pos_encoding = (
            self.w_position_encoder[:w, :].unsqueeze(0).to(input.get_device())
        )
        w_pos_encoding = self.w_linear(w_pos_encoding)  # [1, W, D]

        h_pos_encoding = h_pos_encoding.expand(-1, w, -1).unsqueeze(0)
        w_pos_encoding = w_pos_encoding.expand(
            h, -1, -1).unsqueeze(0)   # h, w, c/2
        # adaptive

        h_pos_encoding = torch.mul(a[:, 0:1, :, :], h_pos_encoding)
        w_pos_encoding = torch.mul(a[:, 1:2, :, :], w_pos_encoding)

        pos_encoding = torch.cat(
            [h_pos_encoding, w_pos_encoding], dim=3)  # [B, H, W, 2*D]

        pos_encoding = pos_encoding.permute(0, 3, 1, 2)  # [B, 2*D, H, W]
        out = input + pos_encoding
        out = self.dropout(out)
        return out

train.py 파일 실행한 결과 정상적으로 동작합니다. 성능은 추후에 피어세션에서 말씀드리겠습니다.