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)