How to define a PixelShuffle Layer?

Question

PixelShuffleLayer come from Real-Time_Single_Image.

Firstly, a $C\times H\times W$ image, after a convolution, now $f^2\cdot C \times H\times W$, then sub-pix, finally get a $ C \times fH\times fW$ image.

Let me take a $3\times 4\times 3$ image with $f=2$ as an example:

SeedRandom[42]
afterConv=RandomReal[1,{2^2*3,4,3}];
ArrayFlatten[Transpose[ArrayReshape[#,{2,2,4,3}],{3,4,1,2}],2]&/@Partition[afterConv,4]

If we do a step decomposition:

MatrixForm/@(m1a=First@Partition[afterConv,4])
MatrixForm[m1b=ArrayReshape[m1a,{2,2,4,3}]]
MatrixForm[m1c=Transpose[m1b,{3,4,1,2}]]
MatrixForm[m1d=ArrayFlatten[m1c,2]]

But NetChain can't receive the functions above, it can only use Layers.

So how do I define a PixelShuffleLayer[f_Integer]?

---

Update-20180920:

Notice that that the color channel, the width and height of the picture are variable.

I tried to decompile the MXNet diagram definition to find how MXNet can do this.

{
    MXNetLink`Symbol`PackagePrivate`mxSymbolCreator["Reshape", {"shape" -> "(0, -4, -1, 4, 0, 0)"}],
    MXNetLink`Symbol`PackagePrivate`mxSymbolCreator["Reshape", {"shape" -> "(0, 0, -4, 2, 2, 0, 0)"}],
    MXNetLink`Symbol`PackagePrivate`mxSymbolCreator["transpose", {"axes" -> "(0, 1, 4, 2, 5, 3)"}],
    MXNetLink`Symbol`PackagePrivate`mxSymbolCreator["Reshape", {"shape" -> "(0, 0, -3, -3)"}]
}

While ReshapeLayer only accept a list of positive integers.

So try NeuralNetworks`MXLayer or NeuralNetworks`DefineLayer.

Answer 1

net = NetChain[
  {
   FlattenLayer[-1],
   TransposeLayer[],
   ReshapeLayer[{7, 7, 3, 3}],
   TransposeLayer[2 <-> 3],
   ReshapeLayer[{21, 21}]
   },
  "Input" -> {9, 7, 7}
  ]

colors = {Red, Green, Blue, Cyan, Magenta, Yellow, Orange, Pink, Purple};

Table[Grid[ConstantArray[i, {7, 7}], Frame -> All, Background -> colors[[i]]], {i, 9}]

net@Table[ConstantArray[i, {7, 7}], {i, 9}] // 
 ArrayPlot[#, ColorRules -> Thread[N@Range[9] -> colors], Mesh -> All] &

Answer 2

I think I have already figure out how NeuralNetworks runs, this should work:

Input: ChannelT[$$InputChannels, TensorT[$$InputSize]]

Output: ChannelT[$$OutputChannels, TensorT[$$OutputSize]]

Parameters:
    $Scaled: PosIntegerT
$$InputChannels: SizeT
$$InputSize: SizeListT[2]
$$OutputChannels: ComputedType[SizeT, $$InputChannels / $Scaled^2]
    $$OutputSize: ComputedType[SizeListT[2], $$InputSize * $Scaled]

Writer: Function[
    input = GetInput["Input", "Batchwise"];
    index = SowNode["reshape", input, "shape" -> {0, -4, -1, #Scaled^2, 0, 0}];
    index = SowNode["reshape", index, "shape" -> {0, 0, -4, #Scaled, #Scaled, 0, 0}];
    index = SowNode["transpose", index, "axes" -> {0, 1, 4, 2, 5, 3}];
    index = SowNode["reshape", index, "shape" -> {0, 0, -3, -3}];
    SetOutput["Output", index]
]

It's not the standard Wolfram Language, it's a kind of DSL(domain-specific language),using special way to import.

Save this as a *.m, then load with:

<< NeuralNetworks`
System`PixelShuffleLayer
file = "the saved file.m";
def = AssociateTo[NeuralNetworks`Private`ReadDefinitionFile[file, "System`"], "Suffix" -> "Layer"]
NeuralNetworks`DefineLayer["PixelShuffle", def];

The test result is the same as @Alexey Golyshev's

net = NetChain[{PixelShuffleLayer[3], FlattenLayer[1]}, "Input" -> {9, 7, 7}]
colors = {Red, Green, Blue, Cyan, Magenta, Yellow, Orange, Pink, Purple};
draw = ArrayPlot[#, ColorRules -> Thread[N@Range[9] -> colors], Mesh -> All]&;
net@Table[ConstantArray[i, {7, 7}], {i, 9}] // draw

The new function PixelShuffleLayer is exactly the same as built-in layers.

It can adapt the shape of the input, just try NetReplacePart[net,"Input"->{27,64,36}].

The only problem is that models with custom layers cannot be imported back.

Unless the exact same NeuralNetworks module is used.

I hope Mathematica 12 can have a better support on the custom layers.