6
$\begingroup$

I am trying to use the structural similarity index, which is defined in terms of a computation over sliding windows, as a neural network loss function.

Obviously one can get a sliding mean using a PoolingLayer, and a convolution layer also computes over sliding windows, but the SSIM formula also calls for a sliding variance and covariance. I'm not sure of the best way to do a general computation (not just pooling or multiplying by a kernel) over sliding windows.

Has anyone tried anything like this? Is there a sensible way to define it in terms of the reshape and NetMapOperator layers?

$\endgroup$
2
$\begingroup$

enter image description here

stddev = NetGraph[
  {
   {ElementwiseLayer[#^2 &], PoolingLayer[{2, 2}, "Function" -> Mean]},
   ThreadingLayer[Sqrt[(#1 - #2^2)*4/3] &]
   },
  {
   NetPort["image"] -> 1 -> 2,
   NetPort["mean"] -> 2
   }
  ]

enter image description here

stddev@<|"image" -> {{{1, 2, 3}, {4, 5, 6}}}, "mean" -> {{{3, 4}}}|>

{{{1.82574,1.82574}}}

StandardDeviation[{1, 2, 4, 5}] // N

1.82574

cov = NetGraph[
  {
   ThreadingLayer[#1*#2 &],
   PoolingLayer[{2, 2}, "Function" -> Mean],
   ThreadingLayer[(#1 - #2*#3)*4/3 &]
   },
  {
   {NetPort["image_1"], NetPort["image_2"]} -> 1 -> 2,
   {2, NetPort["mean_1"], NetPort["mean_2"]} -> 3
   }
  ]

enter image description here

cov@<|"image_1" -> {{{1, 2, 3}, {4, 5, 6}}}, "image_2" -> {{{1, 2, 3}, {4, 5, 6}}}, "mean_1" -> {{{3, 4}}}, "mean_2" -> {{{3, 4}}}|>

{{{3.33333,3.33333}}}

Covariance[{1, 2, 4, 5}, {1, 2, 4, 5}] // N

3.33333

bits = 8;
c1 = (0.01*(2^bits - 1))^2;
c2 = (0.03*(2^bits - 1))^2;

net = NetGraph[
  <|
   "mean_1" -> PoolingLayer[{2, 2}, "Function" -> Mean],
   "mean_2" -> PoolingLayer[{2, 2}, "Function" -> Mean],
   "stddev_1" -> stddev,
   "stddev_2" -> stddev,
   "cov" -> cov,
   "SSIM" -> ThreadingLayer[(2*#1*#2 + c1)*(2*#5 + c2)/((#1^2 + #2^2 + c1)*(#3^2 + #4^2 + c2)) &]
   |>,
  {
   NetPort["image_1"] -> "mean_1",
   NetPort["image_2"] -> "mean_2",
   NetPort["image_1"] -> NetPort["stddev_1", "image"],
   "mean_1" -> NetPort["stddev_1", "mean"],
   NetPort["image_2"] -> NetPort["stddev_2", "image"],
   "mean_2" -> NetPort["stddev_2", "mean"],
   NetPort["image_1"] -> NetPort["cov", "image_1"],
   "mean_1" -> NetPort["cov", "mean_1"],
   NetPort["image_2"] -> NetPort["cov", "image_2"],
   "mean_2" -> NetPort["cov", "mean_2"],
   {"mean_1", "mean_2", "stddev_1", "stddev_2", "cov"} -> "SSIM"
   }
  ]

enter image description here

net@<|"image_1" -> {{{1, 2, 3}, {4, 5, 6}}}, "image_2" -> {{{1, 2, 3}, {4, 5, 6}}}|>

{{{1.,1.}}}

net@<|"image_1" -> {{{1, 2, 3}, {4, 5, 6}}}, "image_2" -> {{{0, 0, 3}, {4, 5, 6}}}|>

{{{0.959683,0.978388}}}

$\endgroup$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.