# # # 0=================================0 # | Kernel Point Convolutions | # 0=================================0 # # # ---------------------------------------------------------------------------------------------------------------------- # # Define network architectures # # ---------------------------------------------------------------------------------------------------------------------- # # Hugues THOMAS - 06/03/2020 # from models.blocks import * import numpy as np class KPCNN(nn.Module): """ Class defining KPCNN """ def __init__(self, config): super(KPCNN, self).__init__() ##################### # Network opperations ##################### # Current radius of convolution and feature dimension layer = 0 r = config.first_subsampling_dl * config.conv_radius in_dim = config.in_features_dim out_dim = config.first_features_dim self.K = config.num_kernel_points # Save all block operations in a list of modules self.block_ops = nn.ModuleList() # Loop over consecutive blocks block_in_layer = 0 for block_i, block in enumerate(config.architecture): # Check equivariance if ('equivariant' in block) and (not out_dim % 3 == 0): raise ValueError('Equivariant block but features dimension is not a factor of 3') # Detect upsampling block to stop if 'upsample' in block: break # Apply the good block function defining tf ops print(block, r) self.block_ops.append(block_decider(block, r, in_dim, out_dim, layer, config)) # Index of block in this layer block_in_layer += 1 # Update dimension of input from output in_dim = out_dim # Detect change to a subsampled layer if 'pool' in block or 'strided' in block: # Update radius and feature dimension for next layer layer += 1 r *= 2 out_dim *= 2 block_in_layer = 0 self.head_mlp = UnaryBlock(out_dim, 1024, False, 0) self.head_softmax = UnaryBlock(1024, config.num_classes, False, 0) ################ # Network Losses ################ self.criterion = torch.nn.CrossEntropyLoss() self.offset_loss = config.offsets_loss self.offset_decay = config.offsets_decay self.output_loss = 0 self.reg_loss = 0 self.l1 = nn.L1Loss() return def forward(self, batch, config): # Save all block operations in a list of modules x = batch.features.clone().detach() # Loop over consecutive blocks for block_op in self.block_ops: x = block_op(x, batch) # Head of network x = self.head_mlp(x, batch) x = self.head_softmax(x, batch) return x def loss(self, outputs, labels): """ Runs the loss on outputs of the model :param outputs: logits :param labels: labels :return: loss """ # TODO: Ignore unclassified points in loss for segmentation architecture # Cross entropy loss self.output_loss = self.criterion(outputs, labels) # Regularization of deformable offsets self.reg_loss = self.offset_regularizer() # Combined loss return self.output_loss + self.reg_loss @staticmethod def accuracy(outputs, labels): """ Computes accuracy of the current batch :param outputs: logits predicted by the network :param labels: labels :return: accuracy value """ predicted = torch.argmax(outputs.data, dim=1) total = labels.size(0) correct = (predicted == labels).sum().item() return correct / total def offset_regularizer(self): fitting_loss = 0 repulsive_loss = 0 for m in self.modules(): if isinstance(m, KPConv) and m.deformable: ############################## # divide offset gradient by 10 ############################## m.unscaled_offsets.register_hook(lambda grad: grad * 0.1) #m.unscaled_offsets.register_hook(lambda grad: print('GRAD2', grad[10, 5, :])) ############## # Fitting loss ############## # Get the distance to closest input point KP_min_d2, _ = torch.min(m.deformed_d2, dim=1) # Normalize KP locations to be independant from layers KP_min_d2 = KP_min_d2 / (m.KP_extent ** 2) # Loss will be the square distance to closest input point. We use L1 because dist is already squared fitting_loss += self.l1(KP_min_d2, torch.zeros_like(KP_min_d2)) ################ # Repulsive loss ################ # Normalized KP locations KP_locs = m.deformed_KP / m.KP_extent # Point should not be close to each other for i in range(self.K): other_KP = torch.cat([KP_locs[:, :i, :], KP_locs[:, i + 1:, :]], dim=1).detach() distances = torch.sqrt(torch.sum((other_KP - KP_locs[:, i:i + 1, :]) ** 2, dim=2)) rep_loss = torch.sum(torch.clamp_max(distances - 1.5, max=0.0) ** 2, dim=1) repulsive_loss += self.l1(rep_loss, torch.zeros_like(rep_loss)) return self.offset_decay * (fitting_loss + repulsive_loss) class KPFCNN(nn.Module): """ Class defining KPFCNN """ def __init__(self, config): super(KPFCNN, self).__init__() ############ # Parameters ############ # Current radius of convolution and feature dimension layer = 0 r = config.first_subsampling_dl * config.conv_radius in_dim = config.in_features_dim out_dim = config.first_features_dim self.K = config.num_kernel_points ##################### # List Encoder blocks ##################### # Save all block operations in a list of modules self.encoder_blocs = nn.ModuleList() self.encoder_skip_dims = [] self.encoder_skips = [] # Loop over consecutive blocks for block_i, block in enumerate(config.architecture): # Check equivariance if ('equivariant' in block) and (not out_dim % 3 == 0): raise ValueError('Equivariant block but features dimension is not a factor of 3') # Detect change to next layer for skip connection if np.any([tmp in block for tmp in ['pool', 'strided', 'upsample', 'global']]): self.encoder_skips.append(block_i) self.encoder_skip_dims.append(in_dim) # Detect upsampling block to stop if 'upsample' in block: break # Apply the good block function defining tf ops self.encoder_blocs.append(block_decider(block, r, in_dim, out_dim, layer, config)) # Update dimension of input from output in_dim = out_dim # Detect change to a subsampled layer if 'pool' in block or 'strided' in block: # Update radius and feature dimension for next layer layer += 1 r *= 2 out_dim *= 2 ##################### # List Decoder blocks ##################### # Save all block operations in a list of modules self.decoder_blocs = nn.ModuleList() self.decoder_concats = [] # Find first upsampling block start_i = 0 for block_i, block in enumerate(config.architecture): if 'upsample' in block: start_i = block_i break # Loop over consecutive blocks for block_i, block in enumerate(config.architecture[start_i:]): # Add dimension of skip connection concat if block_i > 0 and 'upsample' in config.architecture[start_i + block_i - 1]: in_dim += self.encoder_skip_dims[layer] self.decoder_concats.append(block_i) # Apply the good block function defining tf ops self.decoder_blocs.append(block_decider(block, r, in_dim, out_dim, layer, config)) # Update dimension of input from output in_dim = out_dim # Detect change to a subsampled layer if 'upsample' in block: # Update radius and feature dimension for next layer layer -= 1 r *= 0.5 out_dim = out_dim // 2 self.head_mlp = UnaryBlock(out_dim, config.first_features_dim, False, 0) self.head_softmax = UnaryBlock(config.first_features_dim, config.num_classes, False, 0) ################ # Network Losses ################ # Choose segmentation loss if config.segloss_balance == 'none': self.criterion = torch.nn.CrossEntropyLoss() elif config.segloss_balance == 'class': self.criterion = torch.nn.CrossEntropyLoss() elif config.segloss_balance == 'batch': self.criterion = torch.nn.CrossEntropyLoss() else: raise ValueError('Unknown segloss_balance:', config.segloss_balance) self.offset_loss = config.offsets_loss self.offset_decay = config.offsets_decay self.output_loss = 0 self.reg_loss = 0 self.l1 = nn.L1Loss() return def forward(self, batch, config): # Get input features x = batch.features.clone().detach() # Loop over consecutive blocks skip_x = [] for block_i, block_op in enumerate(self.encoder_blocs): if block_i in self.encoder_skips: skip_x.append(x) x = block_op(x, batch) print(block_op) for block_i, block_op in enumerate(self.decoder_blocs): if block_i in self.decoder_concats: x = torch.cat([x, skip_x.pop()], dim=1) x = block_op(x, batch) # Head of network x = self.head_mlp(x, batch) x = self.head_softmax(x, batch) return x def loss(self, outputs, labels): """ Runs the loss on outputs of the model :param outputs: logits :param labels: labels :return: loss """ outputs = torch.transpose(outputs, 0, 1) outputs = outputs.unsqueeze(0) labels = labels.unsqueeze(0) # Cross entropy loss self.output_loss = self.criterion(outputs, labels) # Regularization of deformable offsets self.reg_loss = self.offset_regularizer() # Combined loss return self.output_loss + self.reg_loss @staticmethod def accuracy(outputs, labels): """ Computes accuracy of the current batch :param outputs: logits predicted by the network :param labels: labels :return: accuracy value """ predicted = torch.argmax(outputs.data, dim=1) total = labels.size(0) correct = (predicted == labels).sum().item() return correct / total def offset_regularizer(self): fitting_loss = 0 repulsive_loss = 0 for m in self.modules(): if isinstance(m, KPConv) and m.deformable: ############################## # divide offset gradient by 10 ############################## m.unscaled_offsets.register_hook(lambda grad: grad * 0.1) #m.unscaled_offsets.register_hook(lambda grad: print('GRAD2', grad[10, 5, :])) ############## # Fitting loss ############## # Get the distance to closest input point KP_min_d2, _ = torch.min(m.deformed_d2, dim=1) # Normalize KP locations to be independant from layers KP_min_d2 = KP_min_d2 / (m.KP_extent ** 2) # Loss will be the square distance to closest input point. We use L1 because dist is already squared fitting_loss += self.l1(KP_min_d2, torch.zeros_like(KP_min_d2)) ################ # Repulsive loss ################ # Normalized KP locations KP_locs = m.deformed_KP / m.KP_extent # Point should not be close to each other for i in range(self.K): other_KP = torch.cat([KP_locs[:, :i, :], KP_locs[:, i + 1:, :]], dim=1).detach() distances = torch.sqrt(torch.sum((other_KP - KP_locs[:, i:i + 1, :]) ** 2, dim=2)) rep_loss = torch.sum(torch.clamp_max(distances - 1.5, max=0.0) ** 2, dim=1) repulsive_loss += self.l1(rep_loss, torch.zeros_like(rep_loss)) return self.offset_decay * (fitting_loss + repulsive_loss)