ConvTranspose2D implementation

[nickglazzard1959]

Hello all. I think I've seen some proposals for an implementation of this in MLX but I can't see it in the 0.10.0 release. It looks like the convolution implementation was overhauled recently, but this doesn't seem to have been added. Have I missed it somehow? If not, is there some way of using the existing MLX facilities to implement it (perhaps with some loss of efficiency, but to get the functionality)? Has anyone perhaps already done this? I don't think I have the skills needed to add it to MLX myself, so some higher level approach would be very helpful. I'd be grateful for any ideas.

[awni]

You can use mx.conv_general to implement the transposed convolution (which is the way we would do it if we added a ConvTranspose layer or op. I think there is still room to optimize the underlying implementation but it should not be horribly slow.

This is a useful page to see correspondence between conv and transposed conv.

[francescopisu]

@nickglazzard1959 Do you mind sharing some hints on how to do it using conv_general ? Thanks

[awni]

Here is an example. We should probably add this as a layer.

[francescopisu]

Thanks, much appreciated. It should work with 3d data as well, right ? conv_general docs says only conv1d and conv2d are supported, but I guess it hasn't been updated since MLX added conv3d support.
Anyway, I'll give it a shot later today.

[awni]

Hmm, that is outdated, thanks for catching it. It should work for 3D conv now, but please let us know if you run into any issues.

[nickglazzard1959]

@francescopisu Oh, I'm afraid I haven't done anything yet on this, and I'm not sure when I will, TBH. It looks like the example @awni posted is pretty much what is needed, though.

[petergreis]

I would like to add another vote to add ConvTranspose2D to the mlx baseline. Using it here for a segmentation problem, but it does seem to be running really slow....

[awni]

Could you say more about what you are doing and how you can tell it is slow?

[petergreis]

Sure - the PyTorch version I stopped at 98 epochs and a runtime of 1h:13m. The mlx version has been running 5 hours and is on epoch 66. However, I am profiling the mlx version at the moment. I will do a clean run without the profiler so that we have a fair comparison. The code is doing a segmentation task ala Carvana kaggle challenge, and using an IoU score to judge how good it is.

[petergreis]

And just for more context - I took a PyTorch model that was used for a university assignment (my son at ETH in Zürich) and converted it to mlx. I am attempting to figure out why the IoU score for the mlx version is lower epoch for epoch. Could this be the model initialisation ?

[petergreis]

Blue is mlx, orange is PyTorch