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