import torch import torch.nn as nn import torch.nn.functional as F from models.modelUtils import ChebConv, Pool, residualBlock import torchvision.ops.roi_align as roi_align import numpy as np class EncoderConv(nn.Module): def __init__(self, latents = 64, hw = 32): super(EncoderConv, self).__init__() self.latents = latents self.c = 4 self.size = self.c * np.array([2,4,8,16,32], dtype = np.intc) self.maxpool = nn.MaxPool2d(2) self.dconv_down1 = residualBlock(1, self.size[0]) self.dconv_down2 = residualBlock(self.size[0], self.size[1]) self.dconv_down3 = residualBlock(self.size[1], self.size[2]) self.dconv_down4 = residualBlock(self.size[2], self.size[3]) self.dconv_down5 = residualBlock(self.size[3], self.size[4]) self.dconv_down6 = residualBlock(self.size[4], self.size[4]) self.fc_mu = nn.Linear(in_features=self.size[4]*hw*hw, out_features=self.latents) self.fc_logvar = nn.Linear(in_features=self.size[4]*hw*hw, out_features=self.latents) def forward(self, x): x = self.dconv_down1(x) x = self.maxpool(x) x = self.dconv_down2(x) x = self.maxpool(x) conv3 = self.dconv_down3(x) x = self.maxpool(conv3) conv4 = self.dconv_down4(x) x = self.maxpool(conv4) conv5 = self.dconv_down5(x) x = self.maxpool(conv5) conv6 = self.dconv_down6(x) x = conv6.view(conv6.size(0), -1) # flatten batch of multi-channel feature maps to a batch of feature vectors x_mu = self.fc_mu(x) x_logvar = self.fc_logvar(x) return x_mu, x_logvar, conv6, conv5 class SkipBlock(nn.Module): def __init__(self, in_filters, window): super(SkipBlock, self).__init__() self.window = window self.graphConv_pre = ChebConv(in_filters, 2, 1, bias = False) def lookup(self, pos, layer, salida = (1,1)): B = pos.shape[0] N = pos.shape[1] F = layer.shape[1] h = layer.shape[-1] ## Scale from [0,1] to [0, h] pos = pos * h _x1 = (self.window[0] // 2) * 1.0 _x2 = (self.window[0] // 2 + 1) * 1.0 _y1 = (self.window[1] // 2) * 1.0 _y2 = (self.window[1] // 2 + 1) * 1.0 boxes = [] for batch in range(0, B): x1 = pos[batch,:,0].reshape(-1, 1) - _x1 x2 = pos[batch,:,0].reshape(-1, 1) + _x2 y1 = pos[batch,:,1].reshape(-1, 1) - _y1 y2 = pos[batch,:,1].reshape(-1, 1) + _y2 aux = torch.cat([x1, y1, x2, y2], axis = 1) boxes.append(aux) skip = roi_align(layer, boxes, output_size = salida, aligned=True) vista = skip.view([B, N, -1]) return vista def forward(self, x, adj, conv_layer): pos = self.graphConv_pre(x, adj) skip = self.lookup(pos, conv_layer) return torch.cat((x, skip, pos), axis = 2), pos class Hybrid(nn.Module): def __init__(self, config, downsample_matrices, upsample_matrices, adjacency_matrices): super(Hybrid, self).__init__() self.config = config hw = config['inputsize'] // 32 self.z = config['latents'] self.encoder = EncoderConv(latents = self.z, hw = hw) self.downsample_matrices = downsample_matrices self.upsample_matrices = upsample_matrices self.adjacency_matrices = adjacency_matrices self.kld_weight = 1e-5 n_nodes = config['n_nodes'] self.filters = config['filters'] self.K = 6 self.window = (3,3) # Genero la capa fully connected del decoder outshape = self.filters[-1] * n_nodes[-1] self.dec_lin = torch.nn.Linear(self.z, outshape) self.normalization2u = torch.nn.InstanceNorm1d(self.filters[1]) self.normalization3u = torch.nn.InstanceNorm1d(self.filters[2]) self.normalization4u = torch.nn.InstanceNorm1d(self.filters[3]) self.normalization5u = torch.nn.InstanceNorm1d(self.filters[4]) self.normalization6u = torch.nn.InstanceNorm1d(self.filters[5]) outsize1 = self.encoder.size[4] outsize2 = self.encoder.size[4] # Guardo las capas de convoluciones en grafo self.graphConv_up6 = ChebConv(self.filters[6], self.filters[5], self.K) self.graphConv_up5 = ChebConv(self.filters[5], self.filters[4], self.K) self.SC_1 = SkipBlock(self.filters[4], self.window) self.graphConv_up4 = ChebConv(self.filters[4] + outsize1 + 2, self.filters[3], self.K) self.graphConv_up3 = ChebConv(self.filters[3], self.filters[2], self.K) self.SC_2 = SkipBlock(self.filters[2], self.window) self.graphConv_up2 = ChebConv(self.filters[2] + outsize2 + 2, self.filters[1], self.K) self.graphConv_up1 = ChebConv(self.filters[1], self.filters[0], 1, bias = False) self.pool = Pool() self.reset_parameters() def reset_parameters(self): torch.nn.init.normal_(self.dec_lin.weight, 0, 0.1) def sampling(self, mu, log_var): std = torch.exp(0.5*log_var) eps = torch.randn_like(std) return eps.mul(std).add_(mu) def forward(self, x): self.mu, self.log_var, conv6, conv5 = self.encoder(x) if self.training: z = self.sampling(self.mu, self.log_var) else: z = self.mu x = self.dec_lin(z) x = F.relu(x) x = x.reshape(x.shape[0], -1, self.filters[-1]) x = self.graphConv_up6(x, self.adjacency_matrices[5]._indices()) x = self.normalization6u(x) x = F.relu(x) x = self.graphConv_up5(x, self.adjacency_matrices[4]._indices()) x = self.normalization5u(x) x = F.relu(x) x, pos1 = self.SC_1(x, self.adjacency_matrices[3]._indices(), conv6) x = self.graphConv_up4(x, self.adjacency_matrices[3]._indices()) x = self.normalization4u(x) x = F.relu(x) x = self.pool(x, self.upsample_matrices[0]) x = self.graphConv_up3(x, self.adjacency_matrices[2]._indices()) x = self.normalization3u(x) x = F.relu(x) x, pos2 = self.SC_2(x, self.adjacency_matrices[1]._indices(), conv5) x = self.graphConv_up2(x, self.adjacency_matrices[1]._indices()) x = self.normalization2u(x) x = F.relu(x) x = self.graphConv_up1(x, self.adjacency_matrices[0]._indices()) # Sin relu y sin bias return x, pos1, pos2