Deconvolution with Metal Performance Shaders

Question

Turns out there is no such operation as deconvolution in MPS. The closest analogue in tensorflow is conv2d_transpose.

Is it possible to sort of plug-in custom operations between MPS default operations?

Answer 1

You can write your own Metal compute kernels and execute those in between the MPS operations.

For example:

let commandBuffer = commandQueue.makeCommandBuffer()

. . .

// Do something with an MPSCNN layer:
layer1.encode(commandBuffer: commandBuffer, sourceImage: img1, destinationImage: img2)

// Perform your own compute kernel:
let encoder = commandBuffer.makeComputeCommandEncoder()
encoder.setComputePipelineState(yourOwnComputePipeline)
encoder.setTexture(img2.texture, at: 0)
encoder.setTexture(img3.texture, at: 1)
let threadGroupSize = MTLSizeMake(. . .)
let threadGroups = MTLSizeMake(img2.texture.width / threadGroupSize.width,
                               img2.texture.height / threadGroupSize.height, 1)
encoder.dispatchThreadgroups(threadGroups, threadsPerThreadgroup: threadGroupSize)
encoder.endEncoding()

// Do something with another MPSCNN layer:
layer2.encode(commandBuffer: commandBuffer, sourceImage: img3, destinationImage: img4)

. . .

commandBuffer.commit()

You have to write your own compute kernel in the Metal Shading Language and load this into the yourOwnComputePipeline object. Then you can encode it into the current command buffer whenever you want.

Answer 2

[I am adding this as a new answer because it's a different solution.]

Note that deconvolution in deep learning is also known as "transposed convolution", which means that it's the same as doing a regular convolution but with the kernels horizontally and vertically flipped.

So you should be able to use a regular MPSCNNConvolution layer that takes the MPSImage that you wish to deconvolve as input, and that uses the same kernels as the "forward" convolution step but flipped horizontally and vertically.

The advantage of doing this over writing your own compute kernel is that you can use the very fast kernels from MPS.

Edit: An example. Let's say your conv kernel weights look like this:

1, 2, 3
4, 5, 6
7, 8, 9

Then after flipping the kernel, the weights look like this:

9, 8, 7
6, 5, 4
3, 2, 1

In other words, you need to make a copy of your weights array and reverse it. In memory the first weights look like this:

1, 2, 3, 4, 5, 6, 7, 8, 9

The flipped kernel looks like this in memory, so it's simply the original kernel but in reverse order:

9, 8, 7, 6, 5, 4, 3, 2, 1

Then you make a new convolution layer using that reversed array. This is now your deconv layer.

I don't have Metal sample code to show you, but it's really no different from making a regular MPSCNNConvolution layer. You just have to reverse the weights for the layer.

Answer 3

MPS now provides MPSCNNConvolutionTranspose in macOS X.13 and tvOS/iOS 11.