Exemplo n.º 1
0
    def __init__(self,
                 n_inp=1,
                 feats=[32, 64, 64, 128, 128, 256, 256],
                 blocks=[2, 2, 2, 2, 2, 2, 2],
                 n_pred_p=[1, 1, 1],
                 n_pred_d=[4, 4, 4],
                 L_output=[0, 1, 2],
                 abn=1,
                 dropout_ratio=0.0):
        super().__init__()

        self.relu = nn.ReLU(inplace=True)
        self.pixel_shuffle = nn.PixelShuffle(2)
        self.upsample = nn.Upsample(scale_factor=2, mode='nearest')
        self.n_inp = n_inp
        if abn == 0:
            abnblock = ABN
        elif abn == 1:
            abnblock = InPlaceABN
        else:
            abnblock = InPlaceABNSync

        Conv3d_ABN.define_abn(abnblock)
        Bottleneck_ABN.define_abn(abnblock)

        self.strides = np.array([[1, 1, 1], [2, 2, 2], [4, 4, 4], [8, 8, 8],
                                 [16, 16, 16]])
        for i in range(6):
            if i == 0:
                setattr(self, 'l' + str(i),
                        build_layer(n_inp, feats[i], blocks[i]))
                setattr(self, 'r' + str(i),
                        build_layer(feats[i], feats[i], blocks[i], skip=True))
            else:
                setattr(
                    self, 'l' + str(i),
                    build_layer(feats[i - 1], feats[i], blocks[i], stride=2))
                setattr(self, 'r' + str(i),
                        build_layer(feats[i], feats[i], blocks[i], skip=True))
                setattr(self, 'rt' + str(i),
                        build_layer(feats[i], feats[i - 1], blocks[i]))

        self.r0 = Conv3d_ABN(feats[0], feats[0], kernel_size=3, padding=1)
        self.out_layer = nn.Conv3d(feats[0], 1, kernel_size=1, stride=1)
        self.drop_out = nn.Dropout3d(p=dropout_ratio)
        for c_out1, c_out2, out in zip(n_pred_p, n_pred_d, L_output):
            setattr(self, 'out' + str(out) + 'p',
                    nn.Conv3d(feats[out], c_out1, kernel_size=3, padding=1))
            outd = nn.Conv3d(feats[out], c_out2, kernel_size=3, padding=1)
            #outd.weight.data.fill_(0.0)
            outd.bias.data.fill_(0.0)
            setattr(self, 'out' + str(out) + 'd', outd)

        self.L_output = L_output
        self.strides = self.strides[np.array(L_output)]
Exemplo n.º 2
0
    def __init__(self,
                 n_inp=1,
                 feats=[32, 64, 64, 128, 128, 256, 256],
                 blocks=[2, 2, 2, 2, 2, 2, 2],
                 n_pred_p=[1, 1, 1],
                 n_pred_d=[4, 4, 4],
                 L_output=[0, 1, 2],
                 abn=1,
                 dropout_ratio=0.0):
        """
        :param n_inp: 输入通道数量 [bs, n_inp, z, y, x]
        :param n_pred_p: 预测类别数量 [1,1,1] 二分类
        :param n_pred_d: 预测位置信息 [4,4,4] 四个坐标值
        :param feats: 通道数量 [32, 64, 64, 128, 128, 256, 256]
        :param blocks: residual block的数量 [2,2,2,2,2,2]
        :param L_output: 输出预测结果的层 [0, 1, 2]

        """
        super().__init__()

        self.relu = nn.ReLU(inplace=True)
        self.pixel_shuffle = nn.PixelShuffle(2)
        self.upsample = nn.Upsample(scale_factor=2, mode='nearest')
        self.n_inp = n_inp
        if abn == 0:
            abnblock = ABN
        elif abn == 1:
            abnblock = InPlaceABN
        else:
            abnblock = InPlaceABNSync

        Conv3d_ABN.define_abn(abnblock)
        Bottleneck_ABN.define_abn(abnblock)

        # not used
        self.strides = np.array([[1, 1, 1], [2, 2, 2], [4, 4, 4], [8, 8, 8],
                                 [16, 16, 16]])

        # 6 stages
        for i in range(6):
            # input layer
            if i == 0:
                # 4 blocks
                # 3x3 * 3
                setattr(self, 'l' + str(i),
                        build_layer(n_inp, feats[i], blocks[i]))
                # 3x3 * 2
                setattr(self, 'r' + str(i),
                        build_layer(feats[i], feats[i], blocks[i], skip=True))

            # other layers
            else:
                # 6 blocks
                # stride 2 downsample
                setattr(
                    self, 'l' + str(i),
                    build_layer(feats[i - 1], feats[i], blocks[i], stride=2))
                setattr(self, 'r' + str(i),
                        build_layer(feats[i], feats[i], blocks[i], skip=True))
                # add after upsampled (keep dim with previous block)
                setattr(self, 'rt' + str(i),
                        build_layer(feats[i], feats[i - 1], blocks[i]))

        self.r0 = Conv3d_ABN(feats[0], feats[0], kernel_size=3, padding=1)
        self.out_layer = nn.Conv3d(feats[0], 1, kernel_size=1, stride=1)
        self.drop_out = nn.Dropout3d(p=dropout_ratio)

        # 不同layer的预测类别和位置信息
        for c_out1, c_out2, out in zip(n_pred_p, n_pred_d, L_output):
            # outp = build_layer(feats[out], feats[out] // 4, 0, skip=False)
            outp = nn.Conv3d(feats[out], feats[out], kernel_size=3, padding=1)
            setattr(self, 'before_out' + str(out) + 'p', outp)
            # outd = build_layer(feats[out], feats[out] // 4, 0, skip=False)
            outd = nn.Conv3d(feats[out], feats[out], kernel_size=3, padding=1)
            setattr(self, 'before_out' + str(out) + 'd', outd)

        for c_out1, c_out2, out in zip(n_pred_p, n_pred_d, L_output):
            setattr(self, 'out' + str(out) + 'p',
                    nn.Conv3d(feats[out], c_out1, kernel_size=3, padding=1))
            outd = nn.Conv3d(feats[out], c_out2, kernel_size=3, padding=1)
            # outd.weight.data.fill_(0.0)
            outd.bias.data.fill_(0.0)
            setattr(self, 'out' + str(out) + 'd', outd)

        self.L_output = L_output

        # not used
        self.strides = self.strides[np.array(L_output)]