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)
def load(self, args, strict=True):
if args.data_train == 'ImageNet':
if args.pretrain == 'download' or args.extend == 'download':
state = getattr(models,
'resnet{}'.format(args.depth))(pretrain=True)
elif args.extend:
state = torch.load(args.extend)
else:
common.init_kaiming(self)
return
source = state.state_dict()
target = self.state_dict()
for s, t in zip(source.keys(), target.keys()):
target[t].copy_(source[s])
else:
if args.pretrain:
self.load_state_dict(torch.load(args.pretrain), strict=strict)
def __init__(self,
args,
conv3x3=common.default_conv,
conv1x1=common.default_conv):
super(ResNet_Group, self).__init__()
args = args[0]
self.args = args
m = []
if args.data_train.find('CIFAR') >= 0:
self.expansion = 1
self.n_blocks = (args.depth - 2) // 6
self.in_channels = 16
self.downsample_type = 'A'
n_classes = int(args.data_train[5:])
kwargs = {
'kernel_size': args.kernel_size,
'conv3x3': conv3x3,
}
m.append(common.BasicBlock(args.n_colors, 16, **kwargs))
kwargs['conv3x3'] = BasicBlock
m.append(self.make_layer(16, self.n_blocks, **kwargs))
m.append(self.make_layer(32, self.n_blocks, stride=2, **kwargs))
m.append(self.make_layer(64, self.n_blocks, stride=2, **kwargs))
m.append(nn.AvgPool2d(8))
fc = nn.Linear(64 * self.expansion, n_classes)
elif args.data_train == 'ImageNet':
block_config = {
18: ([2, 2, 2, 2], ResBlock, 1),
34: ([3, 4, 6, 3], ResBlock, 1),
50: ([3, 4, 6, 3], BottleNeck, 4),
101: ([3, 4, 23, 3], BottleNeck, 4),
152: ([3, 8, 36, 3], BottleNeck, 4)
}
n_blocks, self.block, self.expansion = block_config[args.depth]
self.in_channels = 64
self.downsample_type = 'C'
n_classes = 1000
kwargs = {
'conv3x3': conv3x3,
'conv1x1': conv1x1,
}
m.append(
common.BasicBlock(args.n_colors,
64,
7,
stride=2,
conv3x3=conv3x3,
bias=False))
m.append(nn.MaxPool2d(3, 2, padding=1))
m.append(self.make_layer(64, n_blocks[0], 3, **kwargs))
m.append(self.make_layer(128, n_blocks[1], 3, stride=2, **kwargs))
m.append(self.make_layer(256, n_blocks[2], 3, stride=2, **kwargs))
m.append(self.make_layer(512, n_blocks[3], 3, stride=2, **kwargs))
m.append(nn.AvgPool2d(7, 1))
fc = nn.Linear(512 * self.expansion, n_classes)
self.features = nn.Sequential(*m)
self.classifier = fc
# only if when it is child model
if conv3x3 == common.default_conv:
if args.pretrained == 'download' or args.extend == 'download':
state = getattr(models,
'resnet{}'.format(args.depth))(pretrained=True)
elif args.extend:
state = torch.load(args.extend)
else:
common.init_kaiming(self)
return
source = state.state_dict()
target = self.state_dict()
for s, t in zip(source.keys(), target.keys()):
target[t].copy_(source[s])
def __init__(self, args, conv3x3=common.default_conv):
super(ResNet_Basis_Blockwise, self).__init__()
args = args[0]
self.args = args
m = []
if args.data_train.find('CIFAR') >= 0:
self.expansion = 1
self.n_basis1 = args.n_basis1
self.n_basis2 = args.n_basis2
self.n_basis3 = args.n_basis3
self.basis_size1 = args.basis_size1
self.basis_size2 = args.basis_size2
self.basis_size3 = args.basis_size3
self.n_blocks = (args.depth - 2) // 6
self.k_size = args.kernel_size
self.in_channels = 16
self.downsample_type = 'A'
self.basis1 = nn.Parameter(
nn.init.kaiming_uniform_(
torch.Tensor(self.n_basis1, self.basis_size1, self.k_size,
self.k_size)))
self.basis2 = nn.Parameter(
nn.init.kaiming_uniform_(
torch.Tensor(self.n_basis2, self.basis_size2, self.k_size,
self.k_size)))
self.basis3 = nn.Parameter(
nn.init.kaiming_uniform_(
torch.Tensor(self.n_basis3, self.basis_size3, self.k_size,
self.k_size)))
self.block = ResBlockDecom
n_classes = int(args.data_train[5:])
kwargs = {'kernel_size': args.kernel_size, 'conv3x3': conv3x3}
m.append(common.BasicBlock(args.n_colors, 16, **kwargs))
kwargs['conv3x3'] = BasicBlock
kwargs['kernel_size2'] = args.k_size2
kwargs['n_basis'] = self.n_basis1
kwargs['basis_size'] = self.basis_size1
kwargs['basis_l1'] = self.basis1
kwargs['basis_l2'] = self.basis1
m.append(self.make_layer(16, self.n_blocks, **kwargs))
kwargs['n_basis'] = self.n_basis2
kwargs['basis_size'] = self.basis_size2
kwargs['basis_l1'] = self.basis1
kwargs['basis_l2'] = self.basis2
m.append(self.make_layer(32, self.n_blocks, stride=2, **kwargs))
kwargs['n_basis'] = self.n_basis3
kwargs['basis_size'] = self.basis_size3
kwargs['basis_l1'] = self.basis2
kwargs['basis_l2'] = self.basis3
m.append(self.make_layer(64, self.n_blocks, stride=2, **kwargs))
m.append(nn.AvgPool2d(8))
fc = nn.Linear(64 * self.expansion, n_classes)
self.features = nn.Sequential(*m)
self.classifier = fc
# only if when it is child model
if conv3x3 == common.default_conv:
if args.pretrained == 'download' or args.extend == 'download':
state = getattr(models,
'resnet{}'.format(args.depth))(pretrained=True)
elif args.extend:
state = torch.load(args.extend)
else:
common.init_kaiming(self)
return
source = state.state_dict()
target = self.state_dict()
for s, t in zip(source.keys(), target.keys()):
target[t].copy_(source[s])
