[코드 공유] Adaptive 2D Positional Encoding 구현 피드백 요청

[shjas94]

class AdaptivePositionalEncoding2D(nn.Module):
    def __init__(self, in_channels, max_h=64, max_w=128, dropout=0.1):
        super(PositionalEncoding2D, 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)
        #Adaptive
        self.avgpool = nn.AdaptiveAvgPool2d(1)

        self.alpha_1 = nn.Linear(in_channels, in_channels // 2)
        self.relu_a = nn.ReLU(inplace=True)
        self.alpha_2 = nn.Linear(in_channels // 2, in_channels // 2)
        self.sigmoid_a = nn.Sigmoid()

        self.beta_1 = nn.Linear(in_channels, in_channels // 2)
        self.relu_b = nn.ReLU(inplace=True)
        self.beta_2 = nn.Linear(in_channels // 2, in_channels // 2)
        self.sigmoid_b = 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)
        e_a = self.sigmoid_a(self.alpha_2(self.relu_a(self.alpha_1(e))))
        e_b = self.sigmoid_b(self.beta_2(self.relu_b(self.beta_1(e))))
        ########
        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)   # h, w, c/2
        w_pos_encoding = w_pos_encoding.expand(h, -1, -1)   # h, w, c/2

        pos_encoding = torch.cat(
            [torch.mul(h_pos_encoding, e_a), torch.mul(w_pos_encoding, e_b)], dim=2)  # [H, W, 2*D]

        pos_encoding = pos_encoding.permute(2, 0, 1)  # [2*D, H, W]

        out = input + pos_encoding.unsqueeze(0)
        out = self.dropout(out)

        return out

피드백 부탁드립니다
%코드 수정했습니다%

[Atica57]

(추가 요청사항)
starn desnet에 100*400을 input으로 넣었을 경우 내부 레이어 거치면서 output shape은 어떻게 될지 torch summary 결과는?