def ssim_depthwise_convolution(y, t, window_size, stride): n, c, w, h = y.shape window = _create_window(window_size, c, y.xp) mu_y = F.depthwise_convolution_2d(y, window, stride=stride) mu_t = F.depthwise_convolution_2d(t, window, stride=stride) mu_y_sq = F.depthwise_convolution_2d(y * y, window, stride=stride) mu_t_sq = F.depthwise_convolution_2d(t * t, window, stride=stride) mu_ty = F.depthwise_convolution_2d(y * t, window, stride=stride) muy_mut = (mu_y * mu_t) sq_mu_y = mu_y**2 sq_mu_t = mu_t**2 sigma_y_sq = mu_y_sq - sq_mu_y sigma_t_sq = mu_t_sq - sq_mu_t sigma_yt = mu_ty - muy_mut c1 = 0.01**2 c2 = 0.03**2 ssim_map = ((2 * muy_mut + c1) * (2 * sigma_yt + c2)) / ((sq_mu_y + sq_mu_t + c1) * (sigma_y_sq + sigma_t_sq + c2)) return F.mean(ssim_map)
def __call__(self, encs, hiddens, batch_size, prev_image, num_masks, color_channels): """ Learn through StatelessCDNA. Args: encs: An array of computed transformation hiddens: An array of hidden layers batch_size: Size of mini batches prev_image: The image to transform num_masks: Number of masks to apply color_channels: Output color channels Returns: transformed: A list of masks to apply on the previous image """ logger = logging.getLogger(__name__) enc0, enc1, enc2, enc3, enc4, enc5, enc6 = encs hidden1, hidden2, hidden3, hidden4, hidden5, hidden6, hidden7 = hiddens img_height = prev_image.shape[2] img_width = prev_image.shape[3] # CDNA specific enc7 = self.enc7(enc6) enc7 = F.relu(enc7) transformed_list = list([F.sigmoid(enc7)]) # CDNA specific # Predict kernels using linear function of last layer cdna_input = F.reshape(hidden5, (int(batch_size), -1)) cdna_kerns = self.cdna_kerns(cdna_input) # Reshape and normalize # B x C x H x W => B x NUM_MASKS x 1 x H x W cdna_kerns = F.reshape(cdna_kerns, (int(batch_size), self.num_masks, 1, DNA_KERN_SIZE, DNA_KERN_SIZE)) cdna_kerns = F.relu(cdna_kerns - RELU_SHIFT) + RELU_SHIFT norm_factor = F.sum(cdna_kerns, (2, 3, 4), keepdims=True) cdna_kerns = broadcasted_division(cdna_kerns, norm_factor) # Treat the color channel dimension as the batch dimension since the same # transformation is applied to each color channel. # Treat the batch dimension as the channel dimension so that # F.depthwise_convolution_2d can apply a different transformation to each sample. cdna_kerns = F.reshape(cdna_kerns, (int(batch_size), self.num_masks, DNA_KERN_SIZE, DNA_KERN_SIZE)) cdna_kerns = F.transpose(cdna_kerns, (1, 0, 2, 3)) # Swap the batch and channel dimension. prev_image = F.transpose(prev_image, (1, 0, 2, 3)) # Transform the image. transformed = F.depthwise_convolution_2d(prev_image, cdna_kerns, stride=(1, 1), pad=DNA_KERN_SIZE/2) # Transpose the dimensions where they belong. transformed = F.reshape(transformed, (color_channels, int(batch_size), self.num_masks, img_height, img_width)) transformed = F.transpose(transformed, (2, 1, 0, 3, 4)) transformed = F.split_axis(transformed, indices_or_sections=self.num_masks, axis=0) transformed = [F.squeeze(t, axis=0) for t in transformed] transformed_list += transformed return transformed_list, enc7
def forward(self, x): if self.W.array is None: self._initialize_params(x.shape[1]) pad_width = [(0, 0), (0, 0)] + list(map(lambda x: (x, x), self.pad)) x = F.pad(x, pad_width, self.pad_mode) return F.depthwise_convolution_2d(x, self.W, self.b, self.stride, 0)
def check_backward(self, x_data, W_data, b_data, y_grad): args = (x_data, W_data) if b_data is not None: args = args + (b_data,) gradient_check.check_backward( lambda *inputs: functions.depthwise_convolution_2d( *inputs, stride=self.stride, pad=self.pad), args, y_grad, **self.check_backward_options)
def check_forward(self, x_data, W_data, b_data): args1 = (x_data, W_data) args2 = (x_data, W_data) if b_data is not None: args1 = args1 + (b_data,) b_data = sum(numpy.split(b_data, W_data.shape[1])) args2 = args2 + (b_data,) y1 = functions.depthwise_convolution_2d( *args1, stride=self.stride, pad=self.pad) arys = numpy.split(y1.array, self.W.shape[1], axis=1) y1 = sum(arys) y2 = functions.convolution_2d( *args2, stride=self.stride, pad=self.pad).array testing.assert_allclose(y1, y2, **self.check_forward_options)
def __call__(self, x, stride=1, return_f_map=False): if return_f_map: return F.depthwise_convolution_2d(x, self.W, self.b, stride=stride), self.W else: return F.depthwise_convolution_2d(x, self.W, self.b, stride=stride)