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.