kernel size fix

This commit is contained in:
jh 2021-03-13 14:39:36 -08:00
commit 703c790e67

View file

@ -20,7 +20,7 @@ class MsSsim(nn.LayerBase):
y_true_t = tf.transpose(tf.cast(y_true, tf.float32), [0, 2, 3, 1]) y_true_t = tf.transpose(tf.cast(y_true, tf.float32), [0, 2, 3, 1])
y_pred_t = tf.transpose(tf.cast(y_pred, tf.float32), [0, 2, 3, 1]) y_pred_t = tf.transpose(tf.cast(y_pred, tf.float32), [0, 2, 3, 1])
loss = tf.image.ssim_multiscale(y_true_t, y_pred_t, max_val, power_factors=self.power_factors) loss = tf.image.ssim_multiscale(y_true_t, y_pred_t, max_val, power_factors=self.power_factors, filter_size=self.kernel_size)
return (1.0 - loss) / 2.0 return (1.0 - loss) / 2.0
nn.MsSsim = MsSsim nn.MsSsim = MsSsim