class ResNetNNFC1(nn.Module):
def __init__(self, block, num_blocks, num_classes=10, quantizer=87):
super(ResNetNNFC1, self).__init__()
self.in_planes = 64
self.conv1 = nn.Conv2d(3,
64,
kernel_size=3,
stride=1,
padding=1,
bias=False)
self.bn1 = nn.BatchNorm2d(64)
self.layer1 = self._make_layer(block, 64, num_blocks[0], stride=1)
self.layer2 = self._make_layer(block, 128, num_blocks[1], stride=2)
self.layer3 = self._make_layer(block, 256, num_blocks[2], stride=2)
self.layer4 = self._make_layer(block, 512, num_blocks[3], stride=2)
self.linear = nn.Linear(512 * block.expansion, num_classes)
self.timing = False
self.nnfc_compression_layer = CompressionLayer(
encoder_name='nnfc1_encoder',
encoder_params_dict={'quantizer': quantizer},
decoder_name='nnfc1_decoder',
decoder_params_dict={})
def _make_layer(self, block, planes, num_blocks, stride):
strides = [stride] + [1] * (num_blocks - 1)
layers = []
for stride in strides:
layers.append(block(self.in_planes, planes, stride))
self.in_planes = planes * block.expansion
return nn.Sequential(*layers)
def get_compressed_sizes(self):
return self.nnfc_compression_layer.get_compressed_sizes()
def forward(self, x):
out = x
out = F.relu(self.bn1(self.conv1(out)))
if self.timing:
return out
out = self.layer1(out)
out = self.layer2(out)
out = self.nnfc_compression_layer(out)
out = self.layer3(out)
out = self.layer4(out)
out = F.avg_pool2d(out, 4)
out = out.view(out.size(0), -1)
out = self.linear(out)
return out
class Compressor(nn.Module):
def __init__(self, inplanes):
super(Compressor, self).__init__()
# self.compression_layer = CompressionLayer(encoder_name='jpeg_encoder',
# encoder_params_dict={'quantizer' : 36},
# decoder_name='jpeg_decoder',
# decoder_params_dict={})
# self.compression_layer = CompressionLayer(encoder_name='avc_encoder',
# encoder_params_dict={'quantizer' : 42},
# decoder_name='avc_decoder',
# decoder_params_dict={})
self.compression_layer = CompressionLayer(encoder_name='nnfc2_encoder',
encoder_params_dict={},
decoder_name='nnfc2_decoder',
decoder_params_dict={})
# a = torch.arange(0,128).reshape((1, 1, 16, 8)).float()
# print(a.shape, a)
# b = self.compression_layer(a)
# print(b.shape, b)
# print(a == b)
self.sizes = []
self.pad = nn.ReplicationPad2d(2)
# define the bottleneck layers
# expland to 6x the size with 3x3
# mix with 1x1
# bottleneck to same spatial dims, but less channels
#planes = int(inplanes / 2)
# encoder
#t = 12
#self.encoder = LinearBottleneck(inplanes, planes, t=t)
# decoder
#self.decoder = LinearBottleneck(planes, inplanes, t=t)
print('inPlanes', inplanes)
# print('compressPlanes', planes)
# print('t =', t)
def get_compressed_sizes(self):
return self.compression_layer.get_compressed_sizes()
def forward(self, x):
# x = self.encoder(x)
# x_min = float(x.min())
# x_max = float(x.max())
# print(x.max(), x.min())
# thres = 0.5
# x_top = x.clone()
# x_top[x_top <= thres] = 0
# x_bot = x.clone()
# x_bot[x_bot >= -thres] = 0
# x = x_top + x_bot
# density = np.histogram(x.cpu().detach().numpy(), bins=10)
# print(density)
# print(x.max(), x.min())
# visualization code
# print(x.shape)
# d = x[0,:,:,:].cpu().detach().numpy()
# dmin = np.min(d)
# dmax = np.max(d)
# for i in range(x.shape[1]):
# d = x[0,i,:,:].cpu().detach().numpy()
# print(d.shape)
# print(dmin, dmax)
# img = ((255 * (d - dmin)) / (dmax - dmin)).astype(np.uint8)
# imgmin = np.min(img)
# imgmax = np.max(img)
# img = Image.fromarray(img)
# img.save('/home/jemmons/intermediates/intermediate{}.png'.format(i))
# print(x.shape)
x = self.pad(x)
x = self.compression_layer(x)
x = x[:, :, 1:-1, 1:-1]
self.sizes += self.compression_layer.get_compressed_sizes()
# x = self.decoder(x)
# print(np.mean(np.asarray(self.sizes)))
# print(np.median(np.asarray(self.sizes)))
return x