Compute neural network loss over sliding windows?

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?

stddev = NetGraph[
{
{ElementwiseLayer[#^2 &], PoolingLayer[{2, 2}, "Function" -> Mean]},
},
{
NetPort["image"] -> 1 -> 2,
NetPort["mean"] -> 2
}
]


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[
{
PoolingLayer[{2, 2}, "Function" -> Mean],
},
{
{NetPort["image_1"], NetPort["image_2"]} -> 1 -> 2,
{2, NetPort["mean_1"], NetPort["mean_2"]} -> 3
}
]


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"
}
]


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}}}