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)