SoftmaxLayer Incompatible types

Question

Trying to create U-Net similar to this paper from this GitHub project but keep getting "incompatible type" error on the final softmax layer. Any ideas for how to organize data to satisfy Softmax layer? Or is this a bug? Any help is much appreciated. Thanks in advance.

conv[numFilters_Integer] := NetChain[
 {
  ConvolutionLayer[numFilters, {3, 3}, "Stride" -> 1, "PaddingSize" -> 1],
  ElementwiseLayer[Ramp],
  ConvolutionLayer[numFilters, {3, 3}, "Stride" -> 1, "PaddingSize" -> 1],
  ElementwiseLayer[Ramp]
 }
]

net = NetGraph[
 {
  conv[32],
  PoolingLayer[{2, 2}, "Stride" -> 2, "PaddingSize" -> 0],
  conv[64],
  PoolingLayer[{2, 2}, "Stride" -> 2, "PaddingSize" -> 0],
  conv[128],
  NeuralNetworks`UpsampleLayer[2],
  CatenateLayer[],
  conv[64],
  NeuralNetworks`UpsampleLayer[2],
  CatenateLayer[],
  conv[32],
  ConvolutionLayer[2, {1, 1}, "Stride" -> 1, "PaddingSize" -> 0],
  ReshapeLayer[{32*32, 2}],
  SoftmaxLayer[]
  },
 {
  1 -> 2, 2 -> 3, 3 -> 4, 4 -> 5, 5 -> 6,
  6 -> 7, 3 -> 7,
  7 -> 8, 8 -> 9,
  9 -> 10, 1 -> 10,
  10 -> 11, 11 -> 12, 12 -> 13, 13 -> 14
 },
 "Input" -> NetEncoder[{"Image", {32, 32}}]
 ]

Answer 1

Very happy to report that MMA 11.1 has solved this issue with

NetMapOperator[]

which is used together with the SoftmaxLayer[] as the operator to perform the pixel per pixel classification (i.e. segmentation)

Also, new in 11.1 is the ResizeLayer[] which replaces UpsampleLayer[]

So the final network definition is:

net = NetGraph[
    {
        conv[32],
        PoolingLayer[{2, 2}, "Stride" -> 2],
        conv[64],
        PoolingLayer[{2, 2}, "Stride" -> 2],
        conv[128],
        ResizeLayer[{Scaled[2], Scaled[2]}],
        CatenateLayer[],
        conv[64],
        ResizeLayer[{Scaled[2], Scaled[2]}],
        CatenateLayer[],
        conv[32],
        ConvolutionLayer[2, {1, 1}, "Stride" -> 1],
        ReshapeLayer[{32*32, 2}],
        NetMapOperator[SoftmaxLayer[]],
        ReshapeLayer[{32, 32, 2}]
    }, 
    {
        1 -> 2, 2 -> 3, 3 -> 4, 4 -> 5, 5 -> 6, 6 -> 7, 
        3 -> 7, 
        7 -> 8, 8 -> 9, 9 -> 10, 
        1 -> 10, 
        10 -> 11, 11 -> 12, 12 -> 13, 13 -> 14, 14 -> 15
    },
    "Input" -> NetEncoder[{"Image", {32, 32}}]
]

Answer 2

According to documentation, softmax activation in Keras works fine with one more dimenstion (nb_timesteps), it would ignore it anyway.

softmax: Softmax applied across inputs last dimension. Expects shape either (nb_samples, nb_timesteps, nb_dims) or (nb_samples, nb_dims)

Here you need to change the last part a bit:

FlattenLayer[], SoftmaxLayer[], ReshapeLayer[{32 32, 2}]