def __init__(self,
n_feat,
kernel_size,
bias=True,
bn=False,
act=nn.PReLU(),
res_scale=1,
latent_vector=None,
args=None):
super(ResBlock_dhp, self).__init__()
self.finetuning = False
self.res_scale = res_scale
self.layer1 = conv_dhp(n_feat,
n_feat,
kernel_size,
bias=bias,
batchnorm=bn,
act=act,
args=args)
self.layer2 = conv_dhp(n_feat,
n_feat,
kernel_size,
bias=bias,
batchnorm=bn,
act=None,
latent_vector=latent_vector,
args=args)
self.res_scale = res_scale
def __init__(self, args, conv=common.default_conv):
super(SRResNet_DHP, self).__init__(args)
self.width_mult = args.width_mult
self.n_resblock = args.n_resblocks
self.n_feats = int(args.n_feats * self.width_mult)
self.scale = args.scale[0]
# act = nn.PReLU()
self.latent_vector_stage0 = nn.Parameter(torch.randn((3)))
self.latent_vector_stage1 = nn.Parameter(torch.randn((self.n_feats)))
self.head = conv_dhp(self.n_colors,
self.n_feats,
9,
batchnorm=False,
act=nn.PReLU(),
latent_vector=self.latent_vector_stage1,
args=args)
self.body = nn.ModuleList([
ResBlock_dhp(self.n_feats,
3,
bn=True,
act=nn.PReLU(),
latent_vector=self.latent_vector_stage1,
args=args) for _ in range(self.n_resblock)
])
self.body.append(
conv_dhp(self.n_feats,
self.n_feats,
3,
latent_vector=self.latent_vector_stage1,
args=args))
if self.prune_upsampler:
self.tail = nn.ModuleList([
Upsampler(self.scale,
self.n_feats,
act=nn.PReLU(),
latent_vector=self.latent_vector_stage1,
args=args)
])
self.tail.append(
conv_dhp(self.n_feats,
self.n_colors,
3,
batchnorm=False,
args=args))
else:
self.tail = nn.Sequential(*[
common.Upsampler(
conv, self.scale, self.n_feats, act=nn.PReLU()),
conv(self.n_feats, self.n_colors, 3)
])
self.show_latent_vector()
def __init__(self, in_channels, out_channels, args=None):
super(UpConv_dhp, self).__init__()
self.in_channels = in_channels
self.out_channels = out_channels
self.finetuning = False
# self.merge_mode = merge_mode
# self.up_mode = up_mode
# self.flag = flag
# self.upconv = upconv2x2(self.in_channels, self.out_channels, mode=self.up_mode)
self.upconv = conv_dhp(self.in_channels, self.out_channels, 2, stride=2, batchnorm=False, args=args, transpose=True)
self.conv1 = conv_dhp(2 * self.out_channels, self.out_channels, 3, batchnorm=False, act=nn.ReLU(), args=args)
self.conv2 = conv_dhp(self.out_channels, self.out_channels, 3, batchnorm=False, act=nn.ReLU(), args=args)
def __init__(self, in_planes, out_planes, stride=1, embedding_dim=8):
super(DepthwiseSeparableConv_dhp, self).__init__()
self.finetuning = False
self.layer1 = conv_dhp(in_planes,
in_planes,
kernel_size=3,
stride=stride,
groups=in_planes,
embedding_dim=embedding_dim)
self.layer2 = conv_dhp(in_planes,
out_planes,
kernel_size=1,
stride=1,
embedding_dim=embedding_dim)
def __init__(self, in_channels, out_channels, pooling=True, args=None):
super(DownConv_dhp, self).__init__()
self.in_channels = in_channels
self.out_channels = out_channels
self.pooling = pooling
self.finetuning = False
self.conv1 = conv_dhp(self.in_channels, self.out_channels, 3, batchnorm=False, act=nn.ReLU(), args=args)
self.conv2 = conv_dhp(self.out_channels, self.out_channels, 3, batchnorm=False, act=nn.ReLU(), args=args)
# self.conv1 = conv3x3(self.in_channels, self.out_channels)
# self.conv2 = conv3x3(self.out_channels, self.out_channels)
if self.pooling:
self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
def __init__(self, args):
super(DnCNN_DHP, self).__init__(args)
self.width_mult = args.width_mult
n_blocks = args.m_blocks
n_feats = int(args.n_feats * self.width_mult)
kernel_size = 3
# act = nn.ReLU()
self.latent_vector = nn.Parameter(torch.randn((1)))
# define head module
self.head = conv_dhp(args.n_colors, n_feats, kernel_size, batchnorm=False, args=args)
# define body module
self.body = nn.ModuleList([conv_dhp(n_feats, n_feats, kernel_size, act=nn.ReLU(), args=args) for _ in range(n_blocks)])
# define tail module
self.tail = conv_dhp(n_feats, args.n_colors, kernel_size, batchnorm=False, args=args)
self.show_latent_vector()
def __init__(self, args):
super(ResNet_DHP, self).__init__(args=args)
if args.depth <= 56:
self.expansion = 1
self.block = ResBlock_dhp
self.n_blocks = (args.depth - 2) // 6
else:
self.expansion = 4
self.block = BottleNeck_dhp
self.n_blocks = (args.depth - 2) // 9
self.in_channels = 16
self.downsample_type = 'C'
self.prune_same_channels = args.prune_same_channels == 'Yes'
self.latent_vector = nn.Parameter(torch.randn((3)))
stride = 1 if args.data_train.find('CIFAR') >= 0 else 2
self.features = nn.ModuleList([
conv_dhp(args.n_colors,
16,
kernel_size=self.kernel_size,
stride=stride,
embedding_dim=self.embedding_dim)
])
self.features.extend(
self.make_layer(self.n_blocks, 16, self.kernel_size))
self.features.extend(
self.make_layer(self.n_blocks, 32, self.kernel_size, stride=2))
self.features.extend(
self.make_layer(self.n_blocks, 64, self.kernel_size, stride=2))
self.pooling = nn.AvgPool2d(8)
self.classifier = nn.Linear(64 * self.expansion, self.n_classes)
def __init__(self, inp, oup, stride, expand_ratio, latent_vector=None, embedding_dim=8):
super(InvertedResidual_dhp, self).__init__()
self.in_channels = inp
self.out_channels = oup
self.stride = stride
self.finetuning = False
assert stride in [1, 2]
hidden_dim = int(round(inp * expand_ratio))
self.use_res_connect = self.stride == 1 and inp == oup
self.layers = nn.ModuleList()
if expand_ratio != 1:
# pw
self.layers.append(conv_dhp(inp, hidden_dim, 1, embedding_dim=embedding_dim))
self.layers.append(conv_dhp(hidden_dim, hidden_dim, 3, stride=stride, groups=hidden_dim, embedding_dim=embedding_dim))
self.layers.append(conv_dhp(hidden_dim, oup, 1, act=False, latent_vector=latent_vector, embedding_dim=embedding_dim))
def __init__(self, args):
super(ResNet_DHP_Share, self).__init__(args=args)
self.width_mult = args.width_mult
if args.depth <= 56:
self.expansion = 1
self.block = ResBlock_dhp
self.n_blocks = (args.depth - 2) // 6
else:
self.expansion = 4
self.block = BottleNeck_dhp
self.n_blocks = (args.depth - 2) // 9
self.in_channels = int(16 * self.width_mult)
self.downsample_type = 'C'
self.latent_vector_stage0 = nn.Parameter(torch.randn((3)))
self.latent_vector_stage1 = nn.Parameter(
torch.randn((int(16 * self.width_mult) * self.expansion)))
self.latent_vector_stage2 = nn.Parameter(
torch.randn((int(32 * self.width_mult) * self.expansion)))
self.latent_vector_stage3 = nn.Parameter(
torch.randn((int(64 * self.width_mult) * self.expansion)))
stride = 1 if args.data_train.find('CIFAR') >= 0 else 2
v = self.latent_vector_stage1 if self.expansion == 1 else None
self.features = nn.ModuleList([
conv_dhp(args.n_colors,
int(16 * self.width_mult),
kernel_size=self.kernel_size,
stride=stride,
latent_vector=v,
embedding_dim=self.embedding_dim)
])
self.features.extend(
self.make_layer(self.n_blocks,
int(16 * self.width_mult),
self.kernel_size,
latent_vector=self.latent_vector_stage1))
self.features.extend(
self.make_layer(self.n_blocks,
int(32 * self.width_mult),
self.kernel_size,
stride=2,
latent_vector=self.latent_vector_stage2))
self.features.extend(
self.make_layer(self.n_blocks,
int(64 * self.width_mult),
self.kernel_size,
stride=2,
latent_vector=self.latent_vector_stage3))
self.pooling = nn.AvgPool2d(8)
self.classifier = nn.Linear(
int(64 * self.width_mult) * self.expansion, self.n_classes)
self.show_latent_vector()
def __init__(self, args):
super(MobileNetV2_DHP, self).__init__(args=args)
self.width_mult = args.width_mult
self.prune_classifier = args.prune_classifier
if args.data_train == 'ImageNet':
self.cfg = cfg_imagenet
else:
self.cfg = cfg
# NOTE: change conv1 stride 2 -> 1 for CIFAR10
stride = 1 if args.data_train.find('CIFAR') >= 0 else 2
self.latent_vectors = nn.ParameterList([nn.Parameter(torch.randn((3)))] +
[nn.Parameter(torch.randn((int(c[1] * self.width_mult)))) for c in self.cfg])
self.features = nn.ModuleList([conv_dhp(3, int(32 * self.width_mult), kernel_size=3, stride=stride, embedding_dim=self.embedding_dim)])
self.features.extend(self._make_layers(in_planes=int(32 * self.width_mult)))
self.features.append(conv_dhp(int(320 * self.width_mult), int(1280 * self.width_mult), kernel_size=1, stride=1, embedding_dim=self.embedding_dim))
self.classifier = nn.Sequential(nn.Dropout(0.2), nn.Linear(int(1280 * self.width_mult), self.n_classes))
self.show_latent_vector()
def __init__(self,
scale,
n_feat,
bn=False,
act=nn.PReLU(),
bias=True,
latent_vector=None,
args=None):
super(Upsampler, self).__init__()
self.finetuning = False
self.scale = scale
if self.scale == 4: scale = 2
self.upsampler1 = conv_dhp(n_feat,
scale**2 * n_feat,
3,
bias=bias,
batchnorm=bn,
latent_vector=latent_vector,
args=args,
scale=scale)
self.pixel_shuffler1 = nn.PixelShuffle(scale)
if act is not None: self.relu1 = act
if self.scale == 4:
self.upsampler2 = conv_dhp(n_feat,
scale**2 * n_feat,
3,
bias=bias,
batchnorm=bn,
latent_vector=latent_vector,
args=args,
scale=scale)
self.pixel_shuffler2 = nn.PixelShuffle(scale)
if act is not None: self.relu2 = act
else:
raise NotImplementedError(
'Upsampling scale {} is not implemented.'.format(self.scale))
def __init__(self, args):
super(MobileNet_DHP, self).__init__(args=args)
self.width_mult = args.width_mult
self.prune_classifier = args.prune_classifier
self.linear_percentage = args.linear_percentage
self.latent_vector_stage0 = nn.Parameter(torch.randn((3)))
stride = 1 if args.data_train.find('CIFAR') >= 0 else 2
self.features = nn.ModuleList([
conv_dhp(3,
int(32 * self.width_mult),
kernel_size=3,
stride=stride,
embedding_dim=self.embedding_dim)
])
self.features.extend(
self._make_layers(in_planes=int(32 * self.width_mult)))
self.linear = nn.Sequential(
nn.Dropout(0.2),
nn.Linear(int(1024 * self.width_mult), self.n_classes))
self.show_latent_vector()
def __init__(self, args):
super(ResNet_ImageNet_DHP, self).__init__(args=args)
self.width_mult = args.width_mult
if args.depth >= 50:
self.expansion = 4
self.block = BottleNeck_dhp
else:
self.expansion = 1
self.block = ResBlock_dhp
self.linear_percentage = args.linear_percentage
self.config = eval('resnet{}_config'.format(args.depth))
depth = 0
self.depth_cum = []
for d in self.config:
depth += d
self.depth_cum.append(depth)
self.in_channels = int(64 * self.width_mult)
self.latent_vector_stage0 = nn.Parameter(torch.randn((3)))
self.latent_vector_stage1 = nn.Parameter(
torch.randn(int(64 * self.expansion * self.width_mult)))
self.latent_vector_stage2 = nn.Parameter(
torch.randn(int(128 * self.expansion * self.width_mult)))
self.latent_vector_stage3 = nn.Parameter(
torch.randn(int(256 * self.expansion * self.width_mult)))
self.latent_vector_stage4 = nn.Parameter(
torch.randn(int(512 * self.expansion * self.width_mult)))
stride = 2
v = self.latent_vector_stage1 if self.expansion == 1 else None
self.features = nn.ModuleList([
conv_dhp(args.n_colors,
int(64 * self.width_mult),
kernel_size=7,
stride=stride,
latent_vector=v,
embedding_dim=self.embedding_dim),
nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
])
self.features.extend(
self.make_layer(self.config[0],
int(64 * self.width_mult),
self.kernel_size,
latent_vector=self.latent_vector_stage1))
self.features.extend(
self.make_layer(self.config[1],
int(128 * self.width_mult),
self.kernel_size,
stride=2,
latent_vector=self.latent_vector_stage2))
self.features.extend(
self.make_layer(self.config[2],
int(256 * self.width_mult),
self.kernel_size,
stride=2,
latent_vector=self.latent_vector_stage3))
self.features.extend(
self.make_layer(self.config[3],
int(512 * self.width_mult),
self.kernel_size,
stride=2,
latent_vector=self.latent_vector_stage4))
self.pooling = nn.AdaptiveAvgPool2d((1, 1))
self.classifier = nn.Linear(
int(512 * self.expansion * self.width_mult), self.n_classes)
self.show_latent_vector()
def __init__(self, args, conv=common.default_conv):
super(EDSR, self).__init__(args)
width_mult = args.width_mult
n_resblock = args.n_resblocks
n_feats = int(args.n_feats * width_mult)
scale = args.scale[0]
kernel_size = 3
act = nn.ReLU(
) # Pay attention to the difference between inplace and non-inplace operation
rgb_mean = (0.4488, 0.4371, 0.4040)
rgb_std = (1.0, 1.0, 1.0)
self.n_resblock = args.n_resblocks
self.latent_vector_stage0 = nn.Parameter(torch.randn((3)))
self.latent_vector_stage1 = nn.Parameter(torch.randn((n_feats)))
self.sub_mean = common.MeanShift(args.rgb_range, rgb_mean, rgb_std)
# define head module
self.head = conv_dhp(args.n_colors,
n_feats,
kernel_size,
batchnorm=False,
latent_vector=self.latent_vector_stage1,
args=args)
# define body module
self.body = nn.ModuleList([
ResBlock_dhp(n_feats,
kernel_size,
act=act,
res_scale=args.res_scale,
latent_vector=self.latent_vector_stage1,
args=args) for _ in range(n_resblock)
])
self.body.append(
conv_dhp(n_feats,
n_feats,
kernel_size,
batchnorm=False,
latent_vector=self.latent_vector_stage1,
args=args))
# define tail module
if self.prune_upsampler:
self.tail = nn.ModuleList([
Upsampler(scale,
n_feats,
act=None,
latent_vector=self.latent_vector_stage1,
args=args),
conv_dhp(n_feats,
args.n_colors,
kernel_size,
batchnorm=False,
args=args)
])
else:
self.tail = nn.Sequential(*[
common.Upsampler(conv, scale, n_feats),
conv(n_feats, args.n_colors, 3)
])
self.add_mean = common.MeanShift(args.rgb_range, rgb_mean, rgb_std, 1)
self.show_latent_vector()
def __init__(self, args):
"""
Arguments:
in_channels: int, number of channels in the input tensor.
Default is 3 for RGB images.
depth: int, number of MaxPools in the U-Net.
start_filts: int, number of convolutional filters for the
first conv.
up_mode: string, type of upconvolution. Choices: 'transpose'
for transpose convolution or 'upsample' for nearest neighbour
upsampling.
"""
super(UNet_DHP, self).__init__(args)
width_mult = args.width_mult
out_channels = args.n_colors
in_channels = args.n_colors
depth = 5
up_mode='transpose'
start_filts=int(args.n_feats * width_mult)
merge_mode='concat'
if up_mode in ('transpose', 'upsample'):
self.up_mode = up_mode
else:
raise ValueError("\"{}\" is not a valid mode for upsampling. Only \"transpose\" and \"upsample\" are allowed.".format(up_mode))
if merge_mode in ('concat', 'add'):
self.merge_mode = merge_mode
else:
raise ValueError("\"{}\" is not a valid mode for merging up and down paths. Only \"concat\" and \"add\" are allowed.".format(up_mode))
# NOTE: up_mode 'upsample' is incompatible with merge_mode 'add'
if self.up_mode == 'upsample' and self.merge_mode == 'add':
raise ValueError("up_mode \"upsample\" is incompatible with merge_mode \"add\" at the moment because it doesn't make sense to use "
"nearest neighbour to reduce depth channels (by half).")
self.num_classes = out_channels
self.in_channels = in_channels
self.start_filts = start_filts
self.latent_vector = nn.Parameter(torch.randn((1)))
self.down_convs = []
self.up_convs = []
# create the encoder pathway and add to a list
for i in range(depth):
ins = self.in_channels if i == 0 else outs
outs = self.start_filts*(2**i)
pooling = True if i < depth-1 else False
down_conv = DownConv_dhp(ins, outs, pooling=pooling, args=args)
self.down_convs.append(down_conv)
# create the decoder pathway and add to a list
# careful! decoding only requires depth-1 blocks
for i in range(depth-1):
ins = outs
outs = ins // 2
m_flag=True
# m_flag = True if i == (depth - 2) else False
up_conv = UpConv_dhp(ins, outs, args=args)
self.up_convs.append(up_conv)
# add the list of modules to current module
self.down_convs = nn.ModuleList(self.down_convs)
self.up_convs = nn.ModuleList(self.up_convs)
self.conv_final = conv_dhp(outs, self.num_classes, 1, batchnorm=False, args=args)
self.reset_params()
self.show_latent_vector()
