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)
