ablation.py

import numpy as np
import torch
import torch.nn as nn

import ablation_vgg16_c

def crop(data1, data2, crop_h, crop_w):
    _, _, h1, w1 = data1.size()
    _, _, h2, w2 = data2.size()
    assert(h2 <= h1 and w2 <= w1)
    data = data1[:, :, crop_h:crop_h+h2, crop_w:crop_w+w2]
    return data

def get_upsampling_weight(in_channels, out_channels, kernel_size):
    """Make a 2D bilinear kernel suitable for upsampling"""
    factor = (kernel_size + 1) // 2
    if kernel_size % 2 == 1:
        center = factor - 1
    else:
        center = factor - 0.5
    og = np.ogrid[:kernel_size, :kernel_size]
    filt = (1 - abs(og[0] - center) / factor) * \
           (1 - abs(og[1] - center) / factor)
    weight = np.zeros((in_channels, out_channels, kernel_size, kernel_size),
                      dtype=np.float64)
    weight[range(in_channels), range(out_channels), :, :] = filt
    return torch.from_numpy(weight).float()

class MSBlock(nn.Module):
    def __init__(self, c_in, k=3, rate=4):
        super(MSBlock, self).__init__()
        c_out = c_in
        self.k = k
        self.rate = rate

        self.conv = nn.Conv2d(c_in, 32, 3, stride=1, padding=1)
        self.relu = nn.ReLU(inplace=True)

        if k>=1:
            dilation = self.rate*1 if self.rate >= 1 else 1
            self.conv1 = nn.Conv2d(32, 32, 3, stride=1, dilation=dilation, padding=dilation)
            self.relu1 = nn.ReLU(inplace=True)
        if k>=2:
            dilation = self.rate*2 if self.rate >= 1 else 1
            self.conv2 = nn.Conv2d(32, 32, 3, stride=1, dilation=dilation, padding=dilation)
            self.relu2 = nn.ReLU(inplace=True)
        if k>=3:
            dilation = self.rate*3 if self.rate >= 1 else 1
            self.conv3 = nn.Conv2d(32, 32, 3, stride=1, dilation=dilation, padding=dilation)
            self.relu3 = nn.ReLU(inplace=True)
        if k>=4:
            dilation = self.rate*4 if self.rate >= 1 else 1
            self.conv4 = nn.Conv2d(32, 32, 3, stride=1, dilation=dilation, padding=dilation)
            self.relu4 = nn.ReLU(inplace=True)

        self._initialize_weights()

    def forward(self, x):
        o = self.relu(self.conv(x))
        if self.k>=1:
            o1 = self.relu1(self.conv1(o))
        if self.k>=2:
            o2 = self.relu2(self.conv2(o))
        if self.k>=3:
            o3 = self.relu3(self.conv3(o))
        if self.k>=4:
            o4 = self.relu4(self.conv4(o))
        if self.k < 1:
            return o
        elif self.k>=4:
            return o+o1+o2+o3+o4
        elif self.k>=3:
            return o + o1 + o2 + o3
        elif self.k>=2:
            return o + o1 + o2
        elif self.k>=1:
            return o + o1

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                m.weight.data.normal_(0, 0.01)
                if m.bias is not None:
                    m.bias.data.zero_()


class BDCN(nn.Module):
    def __init__(self, pretrain=None, logger=None, ms=True, block=5, bdcn=True, direction='both', k=3, rate=4):
        super(BDCN, self).__init__()
        if logger:
            logger.info(ms)
            logger.info(block)
            logger.info(bdcn)
        self.pretrain = pretrain
        self.ms = ms
        self.block = block
        self.bdcn = bdcn
        self.dir = direction
        self.k = k
        t = 1

        self.features = ablation_vgg16_c.VGG16_C(pretrain, logger, block=block)
        if ms:
            self.msblock1_1 = MSBlock(64, k, rate)
            self.msblock1_2 = MSBlock(64, k, rate)
        else:
            t = 2
        self.conv1_1_down = nn.Conv2d(32*t, 21, (1, 1), stride=1)
        self.conv1_2_down = nn.Conv2d(32*t, 21, (1, 1), stride=1)
        self.score_dsn1 = nn.Conv2d(21, 1, (1, 1), stride=1)
        self.score_dsn1_1 = nn.Conv2d(21, 1, 1, stride=1)
        if block >= 2:
            if ms:
                self.msblock2_1 = MSBlock(128, k, rate)
                self.msblock2_2 = MSBlock(128, k, rate)
            else:
                t = 4
            self.conv2_1_down = nn.Conv2d(32*t, 21, (1, 1), stride=1)
            self.conv2_2_down = nn.Conv2d(32*t, 21, (1, 1), stride=1)
            self.score_dsn2 = nn.Conv2d(21, 1, (1, 1), stride=1)
            self.score_dsn2_1 = nn.Conv2d(21, 1, (1, 1), stride=1)
            self.upsample_2 = nn.ConvTranspose2d(1, 1, 4, stride=2, bias=False)
        if block >= 3:
            if ms:
                self.msblock3_1 = MSBlock(256, k, rate)
                self.msblock3_2 = MSBlock(256, k, rate)
                self.msblock3_3 = MSBlock(256, k, rate)
            else:
                t = 8
            self.conv3_1_down = nn.Conv2d(32*t, 21, (1, 1), stride=1)
            self.conv3_2_down = nn.Conv2d(32*t, 21, (1, 1), stride=1)
            self.conv3_3_down = nn.Conv2d(32*t, 21, (1, 1), stride=1)
            self.score_dsn3 = nn.Conv2d(21, 1, (1, 1), stride=1)
            self.score_dsn3_1 = nn.Conv2d(21, 1, (1, 1), stride=1)
            self.upsample_4 = nn.ConvTranspose2d(1, 1, 8, stride=4, bias=False)
        if block >= 4:
            if ms:
                self.msblock4_1 = MSBlock(512, k, rate)
                self.msblock4_2 = MSBlock(512, k, rate)
                self.msblock4_3 = MSBlock(512, k, rate)
            else:
                t = 16
            self.conv4_1_down = nn.Conv2d(32*t, 21, (1, 1), stride=1)
            self.conv4_2_down = nn.Conv2d(32*t, 21, (1, 1), stride=1)
            self.conv4_3_down = nn.Conv2d(32*t, 21, (1, 1), stride=1)
            self.score_dsn4 = nn.Conv2d(21, 1, (1, 1), stride=1)
            self.score_dsn4_1 = nn.Conv2d(21, 1, (1, 1), stride=1)
            self.upsample_8 = nn.ConvTranspose2d(1, 1, 16, stride=8, bias=False)
        if block >=5:
            if ms:
                self.msblock5_1 = MSBlock(512, k, rate)
                self.msblock5_2 = MSBlock(512, k, rate)
                self.msblock5_3 = MSBlock(512, k, rate)
            else:
                t = 16
            self.conv5_1_down = nn.Conv2d(32*t, 21, (1, 1), stride=1)
            self.conv5_2_down = nn.Conv2d(32*t, 21, (1, 1), stride=1)
            self.conv5_3_down = nn.Conv2d(32*t, 21, (1, 1), stride=1)
            self.score_dsn5 = nn.Conv2d(21, 1, (1, 1), stride=1)
            self.score_dsn5_1 = nn.Conv2d(21, 1, (1, 1), stride=1)
            self.upsample_8_5 = nn.ConvTranspose2d(1, 1, 16, stride=8, bias=False)
        if bdcn and self.dir == 'both':
            c = block * 2
        else:
            c = block
        self.fuse = nn.Conv2d(c, 1, 1, stride=1)

        self._initialize_weights(logger)

    def forward(self, x):
        features = self.features(x)
        if self.ms:
            sum1 = self.conv1_1_down(self.msblock1_1(features[0])) + \
                    self.conv1_2_down(self.msblock1_2(features[1]))
        else:
            sum1 = self.conv1_1_down(features[0]) + \
                    self.conv1_2_down(features[1])
        s1 = self.score_dsn1(sum1)
        if self.bdcn:
            s11 = self.score_dsn1_1(sum1)
        if self.block >= 2:
            if self.ms:
                sum2 = self.conv2_1_down(self.msblock2_1(features[2])) + \
                    self.conv2_2_down(self.msblock2_2(features[3]))
            else:
                sum2 = self.conv2_1_down(features[2]) + \
                    self.conv2_2_down(features[3])
            s2 = self.score_dsn2(sum2)
            s2 = self.upsample_2(s2)
            s2 = crop(s2, x, 1, 1)
            if self.bdcn:
                s21 = self.score_dsn2_1(sum2)
                s21 = self.upsample_2(s21)
                s21 = crop(s21, x, 1, 1)
        if self.block >= 3:
            if self.ms:
                sum3 = self.conv3_1_down(self.msblock3_1(features[4])) + \
                    self.conv3_2_down(self.msblock3_2(features[5])) + \
                    self.conv3_3_down(self.msblock3_3(features[6]))
            else:
                sum3 = self.conv3_1_down(features[4]) + \
                    self.conv3_2_down(features[5]) + \
                    self.conv3_3_down(features[6])
            s3 = self.score_dsn3(sum3)
            s3 =self.upsample_4(s3)
            s3 = crop(s3, x, 2, 2)
            if self.bdcn:
                s31 = self.score_dsn3_1(sum3)
                s31 =self.upsample_4(s31)
                s31 = crop(s31, x, 2, 2)
        if self.block >= 4:
            if self.ms:
                sum4 = self.conv4_1_down(self.msblock4_1(features[7])) + \
                    self.conv4_2_down(self.msblock4_2(features[8])) + \
                    self.conv4_3_down(self.msblock4_3(features[9]))
            else:
                sum4 = self.conv4_1_down(features[7]) + \
                    self.conv4_2_down(features[8]) + \
                    self.conv4_3_down(features[9])
            s4 = self.score_dsn4(sum4)
            s4 = self.upsample_8(s4)
            s4 = crop(s4, x, 4, 4)
            if self.bdcn:
                s41 = self.score_dsn4_1(sum4)
                s41 = self.upsample_8(s41)
                s41 = crop(s41, x, 4, 4)
        if self.block >= 5:
            if self.ms:
                sum5 = self.conv5_1_down(self.msblock5_1(features[10])) + \
                    self.conv5_2_down(self.msblock5_2(features[11])) + \
                    self.conv5_3_down(self.msblock5_3(features[12]))
            else:
                sum5 = self.conv5_1_down(features[10]) + \
                    self.conv5_2_down(features[11]) + \
                    self.conv5_3_down(features[12])
            s5 = self.score_dsn5(sum5)
            s5 = self.upsample_8_5(s5)
            s5 = crop(s5, x, 0, 0)
            if self.bdcn:
                s51 = self.score_dsn5_1(sum5)
                s51 = self.upsample_8_5(s51)
                s51 = crop(s51, x, 0, 0)
        if self.bdcn:
            if self.block >= 5:
                o1, o2, o3, o4, o5 = s1.detach(), s2.detach(), s3.detach(), s4.detach(), s5.detach()
                o11, o21, o31, o41, o51 = s11.detach(), s21.detach(), s31.detach(), s41.detach(), s51.detach()
                p1_1 = s1
                p2_1 = s2 + o1
                p3_1 = s3 + o2 + o1
                p4_1 = s4 + o3 + o2 + o1
                p5_1 = s5 + o4 + o3 + o2 + o1
                p1_2 = s11 + o21 + o31 + o41 + o51
                p2_2 = s21 + o31 + o41 + o51
                p3_2 = s31 + o41 + o51
                p4_2 = s41 + o51
                p5_2 = s51
                if self.dir == 'both':
                    fuse = self.fuse(torch.cat([p1_1, p2_1, p3_1, p4_1, p5_1, p1_2, p2_2, p3_2, p4_2, p5_2], 1))
                    return [p1_1, p2_1, p3_1, p4_1, p5_1, p1_2, p2_2, p3_2, p4_2, p5_2, fuse]
                if self.dir == 'd2s':
                    fuse = self.fuse(torch.cat([p1_1, p2_1, p3_1, p4_1, p5_1], 1))
                    return [p1_1, p2_1, p3_1, p4_1, p5_1, fuse]
                elif self.dir == 's2d':
                    fuse = self.fuse(torch.cat([p1_2, p2_2, p3_2, p4_2, p5_2], 1))
                    return [p1_2, p2_2, p3_2, p4_2, p5_2, fuse]
            elif self.block >= 4:
                o1, o2, o3, o4 = s1.detach(), s2.detach(), s3.detach(), s4.detach()
                o11, o21, o31, o41 = s11.detach(), s21.detach(), s31.detach(), s41.detach()
                p1_1 = s1
                p2_1 = s2 + o1
                p3_1 = s3 + o2 + o1
                p4_1 = s4 + o3 + o2 + o1
                p1_2 = s11 + o21 + o31 + o41
                p2_2 = s21 + o31 + o41
                p3_2 = s31 + o41
                p4_2 = s41
                fuse = self.fuse(torch.cat([p1_1, p2_1, p3_1, p4_1,p1_2, p2_2, p3_2, p4_2], 1))
                return [p1_1, p2_1, p3_1, p4_1, p1_2, p2_2, p3_2, p4_2, fuse]
            elif self.block >= 3:
                o1, o2, o3 = s1.detach(), s2.detach(), s3.detach()
                o11, o21, o31 = s11.detach(), s21.detach(), s31.detach()
                p1_1 = s1
                p2_1 = s2 + o1
                p3_1 = s3 + o2 + o1
                p1_2 = s11 + o21 + o31
                p2_2 = s21 + o31
                p3_2 = s31
                fuse = self.fuse(torch.cat([p1_1, p2_1, p3_1, p1_2, p2_2, p3_2], 1))
                return [p1_1, p2_1, p3_1, p1_2, p2_2, p3_2, fuse]
            elif self.block >= 2:
                o1, o2 = s1.detach(), s2.detach()
                o11, o21 = s11.detach(), s21.detach()
                p1_1 = s1
                p2_1 = s2 + o1
                p1_2 = s11 + o21
                p2_2 = s21
                fuse = self.fuse(torch.cat([p1_1, p2_1, p1_2, p2_2], 1))
                return [p1_1, p2_1, p1_2, p2_2, fuse]

        concat = s1
        res = [s1]
        if self.block >= 2:
            concat = torch.cat([concat, s2], 1)
            res = [s1, s2]
        if self.block >= 3:
            concat = torch.cat([concat, s3], 1)
            res = [s1, s2, s3]
        if self.block >= 4:
            concat = torch.cat([concat, s4], 1)
            res = [s1, s2, s3, s4]
        if self.block >= 5:
            concat = torch.cat([concat, s5], 1)
            res = [s1, s2, s3, s4, s5]
        fuse = self.fuse(concat)
        res.append(fuse)
        return res

    def _initialize_weights(self, logger=None):
        for name, param in self.state_dict().items():
            if self.pretrain and 'features' in name:
                continue
            elif 'down' in name:
                param.zero_()
            elif 'upsample' in name:
                if logger:
                    logger.info('init upsamle layer %s ' % name)
                k = int(name.split('.')[0].split('_')[1])
                param.copy_(get_upsampling_weight(1, 1, k*2))
            elif 'fuse' in name:
                if logger:
                    logger.info('init params %s ' % name)
                if 'bias' in name:
                    param.zero_()
                else:
                    nn.init.constant(param, 0.080)
            else:
                if logger:
                    logger.info('init params %s ' % name)
                if 'bias' in name:
                    param.zero_()
                else:
                    param.normal_(0, 0.01)

if __name__ == '__main__':
    model = BDCN('./caffemodel2pytorch/vgg16.pth')
    for name, param in model.state_dict().items():
        print name, param