[코드 공유] 백본 교체 코드 공유

[shjas94]

class Backbone(nn.Module):
    def __init__(self, input_channel, num_in_features, output_channel=256, dropout_rate=0.2, depth=16, growth_rate=24):
        super(Backbone, self).__init__()
        num_features = num_in_features
        num_features = num_features + depth * growth_rate
        num_features = num_features // 2
        num_features = num_features + depth * growth_rate

        self.prev = nn.Conv2d(input_channel, 3, (1, 1), bias=False)
        effnet = timm.create_model('tf_efficientnet_b5_ns', pretrained=True)
        self.head = nn.Sequential(effnet.conv_stem, effnet.bn1, effnet.act1)
        block = list(list(effnet.blocks.named_children())[0][1])
        block_01 = list(list(effnet.blocks.named_children())[1][1])
        block = block + block_01
        self.blocks = nn.Sequential(*block)
        self.maxpool = nn.MaxPool2d(kernel_size=(2, 2), stride=2)
        self.conv_last = nn.Conv2d(
            40, num_features // 2, 3, stride=1, padding=1)

    def forward(self, x):
        out = self.prev(x)
        out = self.head(out)
        out = self.blocks(out)
        out = self.maxpool(out)
        return F.relu(self.conv_last(out), inplace=True)

변경된 백본 공유합니다. output feature map은 동일하니 그대로 넣고 Argument만 변경해주시면 됩니다.