def get_prune_list(resnet_channel_list, prob_list, dcfg, expand_rate=0.0001):
import numpy as np
prune_list = []
idx = 0
chn_input_full, chn_output_full = 3, resnet_channel_list[0]
dnas_conv = lambda input, output: DNAS.Conv2d(input, output, 1, 1, 1, False, dcfg=dcfg)
conv = dnas_conv(chn_input_full, chn_output_full)
chn_output_prune = int(np.round(
min(torch.dot(prob_list[idx], conv.out_plane_list).item(), chn_output_full)
))
chn_output_prune += int(np.ceil(expand_rate*(chn_output_full-chn_output_prune)))
prune_list.append(chn_output_prune)
chn_input_full = chn_output_full
idx += 1
for blocks in resnet_channel_list[1:]:
blocks_list = []
for block in blocks:
block_prune_list = []
for chn_output_full in block:
conv = DNAS.Conv2d(chn_input_full, chn_output_full, 1, 1, 1, False, dcfg=dcfg)
print(prob_list[idx], conv.out_plane_list, torch.dot(prob_list[idx], conv.out_plane_list).item())
chn_output_prune = int(np.round(
min(torch.dot(prob_list[idx], conv.out_plane_list).item(), chn_output_full)
))
chn_output_prune += int(np.ceil(expand_rate*(chn_output_full-chn_output_prune)))
block_prune_list.append(chn_output_prune)
chn_input_full = chn_output_full
idx += 1
blocks_list.append(block_prune_list)
prune_list.append(blocks_list)
return prune_list
def __init__(self,
in_planes,
out_planes_list,
stride=2,
project=False,
dcfg=None):
super(BottleneckGated, self).__init__()
out_planes_1 = out_planes_list[0]
out_planes_2 = out_planes_list[1]
out_planes_3 = out_planes_list[2]
assert dcfg is not None
self.dcfg = dcfg
self.dcfg_nonreuse = dcfg.copy()
self.bn0 = nn.BatchNorm2d(in_planes,
momentum=_BATCH_NORM_DECAY,
eps=_EPSILON)
self.conv1 = DNAS.Conv2d(in_planes,
out_planes_1,
kernel_size=1,
stride=1,
padding=0,
bias=False,
dcfg=self.dcfg_nonreuse)
self.bn1 = nn.BatchNorm2d(out_planes_1,
momentum=_BATCH_NORM_DECAY,
eps=_EPSILON)
self.conv2 = DNAS.Conv2d(out_planes_1,
out_planes_2,
kernel_size=3,
stride=stride,
padding=1,
bias=False,
dcfg=self.dcfg_nonreuse)
self.bn2 = nn.BatchNorm2d(out_planes_2,
momentum=_BATCH_NORM_DECAY,
eps=_EPSILON)
self.shortcut = None
# if stride != 1 and len(out_planes_list) > 2:
if project:
self.shortcut = DNAS.Conv2d(in_planes,
out_planes_list[-1],
kernel_size=1,
stride=stride,
padding=0,
bias=False,
dcfg=self.dcfg_nonreuse)
self.dcfg.reuse_gate = self.shortcut.gate
self.conv3 = DNAS.Conv2d(out_planes_2,
out_planes_3,
kernel_size=1,
stride=1,
padding=0,
bias=False,
dcfg=self.dcfg)
def __init__(self,
num_blocks,
num_planes,
num_colors=3,
scale=1,
res_scale=0.1):
super(EDSRGated, self).__init__()
self.num_blocks = num_blocks
self.act = nn.ReLU(inplace=True)
self.sub_mean = MeanShift(1)
self.add_mean = MeanShift(1, sign=1)
self.conv0 = dnas.Conv2d(num_colors,
num_planes,
3,
stride=1,
padding=1,
bias=False,
dcfg=self.dcfg)
self.dcfg.reuse_gate = self.conv0.gate
self.blocks = list()
for _ in range(num_blocks):
self.blocks.append(EDSRBlockGated(num_planes, res_scale,
self.dcfg))
m_tail = list()
m_tail.append(Upsampler(scale, num_planes))
m_tail.append(
nn.Conv2d(in_channels=num_planes,
out_channels=num_colors,
kernel_size=3,
stride=1,
padding=1,
bias=False))
self.tail = nn.Sequential(*m_tail)
def __init__(self, n_layer, n_class, channel_lists, dcfg):
super(ResNet, self).__init__()
if n_layer not in cfg.keys():
print('Numer of layers Error: ', n_layer)
exit(1)
self.n_class = n_class
self.channel_lists = channel_lists
self.block_n_cell = cfg[n_layer]
self.dcfg = dcfg
self.base_n_channel = channel_lists[0]
self.imagenet = len(self.block_n_cell) > 3
if self.imagenet:
self.cell_fn = BottleneckGated if n_layer >= 50 else ResidualBlockGated
self.conv0 = DNAS.Conv2d(3,
self.base_n_channel,
7,
stride=2,
padding=3,
bias=False,
dcfg=self.dcfg)
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
else:
self.cell_fn = ResidualBlockGated
self.conv0 = DNAS.Conv2d(3,
self.base_n_channel,
3,
stride=1,
padding=1,
bias=False,
dcfg=self.dcfg)
self.dcfg.reuse_gate = self.conv0.gate
self.block_list = self._block_layers()
self.bn = nn.BatchNorm2d(channel_lists[-1][-1][-1],
momentum=_BATCH_NORM_DECAY,
eps=_EPSILON)
self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
self.fc = DNAS.Linear(channel_lists[-1][-1][-1],
self.n_class,
dcfg=self.dcfg)
self.apply(_weights_init)
def __init__(self, num_planes, res_scale=1, dcfg=None):
super(EDSRBlockGated, self).__init__()
assert dcfg is not None
self.dcfg = dcfg
self.dcfg_nonreuse = dcfg.copy()
self.conv1 = dnas.Conv2d(num_planes,
num_planes,
kernel_size=3,
stride=1,
padding=1,
bias=False,
dcfg=self.dcfg_nonreuse)
self.act = nn.ReLU(inplace=True)
self.conv2 = dnas.Conv2d(num_planes,
num_planes,
kernel_size=3,
stride=1,
padding=1,
bias=False,
dcfg=self.dcfg)
self.res_scale = res_scale