Exemplo n.º 1
0
    def __init__(self, args, conv3x3=common.default_conv, conv1x1=None):
        super(DenseNet_Basis, self).__init__()
        # from IPython import embed; embed()
        args = args[0]
        n_basis = args.n_basis
        n_group = args.n_group
        kernel_size = args.k_size1

        n_blocks = (args.depth - 4) // 3
        if args.bottleneck: n_blocks //= 2

        k = args.k
        c_in = 2 * k

        basis_size = k * (n_blocks * 3 + 1) // n_group
        transition_group = args.transition_group
        def _dense_block(basis, in_channels):
            module_list = []
            for _ in range(n_blocks):
                if args.bottleneck:
                    module_list.append(BottleNeck(in_channels, k, conv=conv3x3))
                else:
                    module_list.append(Dense(basis, in_channels, k, conv=BasicBlock, args=args))
                in_channels += k

            return nn.Sequential(*module_list)

        self.basis = nn.Parameter(nn.init.kaiming_uniform_(torch.Tensor(n_basis, basis_size, kernel_size, kernel_size)))#.to(self.device))
        module_list = []
        module_list.append(conv3x3(args.n_colors, c_in, 3, bias=False))

        for i in range(3):
            module_list.append(_dense_block(self.basis, c_in))
            c_in += k * n_blocks
            if i < 2:
                c_out = int(math.floor(args.reduction * c_in))
                module_list.append(Transition(c_in, transition_group, c_out))
                c_in = c_out

        module_list.append(common.default_norm(c_in))
        module_list.append(common.default_act())
        module_list.append(nn.AvgPool2d(8))
        self.features = nn.Sequential(*module_list)

        if args.data_train == 'ImageNet':
            n_classes = 1000
        else:
            if args.data_train.find('CIFAR') >= 0:
                n_classes = int(args.data_train[5:])

        self.classifier = nn.Linear(c_in, n_classes)

        common.init_kaiming(self)
    def __init__(self, args, conv3x3=common.default_conv, conv1x1=None):
        super(DenseNet_Group, self).__init__()
        args = args[0]
        n_blocks = (args.depth - 4) // 3
        if args.bottleneck: n_blocks //= 2

        k = args.k
        c_in = 2 * k

        def _dense_block(in_channels):
            module_list = []
            for _ in range(n_blocks):
                if args.bottleneck:
                    module_list.append(BottleNeck(in_channels, k,
                                                  conv=conv3x3))
                else:
                    group_size = 3 if in_channels <= 252 else args.group_size
                    module_list.append(
                        Dense(in_channels,
                              k,
                              group_size=group_size,
                              conv=conv3x3))
                in_channels += k

            return nn.Sequential(*module_list)

        module_list = []
        module_list.append(conv3x3(args.n_colors, c_in, 3, bias=False))

        for i in range(3):
            module_list.append(_dense_block(c_in))
            c_in += k * n_blocks
            if i < 2:
                c_out = int(math.floor(args.reduction * c_in))
                module_list.append(Transition(c_in, c_out))
                c_in = c_out

        module_list.append(common.default_norm(c_in))
        module_list.append(common.default_act())
        module_list.append(nn.AvgPool2d(8))
        self.features = nn.Sequential(*module_list)

        if args.data_train == 'ImageNet':
            n_classes = 1000
        else:
            if args.data_train.find('CIFAR') >= 0:
                n_classes = int(args.data_train[5:])

        self.classifier = nn.Linear(c_in, n_classes)

        common.init_kaiming(self)