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....

[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

[awni]

First, that's a lot slower. It would be useful to know the results of your profile. Or if you have some reproducible benchmarks to share that would also be helpful.

Second, as for the score being lower in MLX. That I do not know. There could be an error in the conversion in any number of places. Or the behavior could be different in some places in MLX than in Torch (including model initialization like you said or default settings in optimizers or other places). If you are in fp32 and the computation is matched, the two should not be that different.

A way to debug that is to try and load the same weights and make sure you get the same scores as a first step. That will at least make sure the forward computation is the same.

[petergreis]

Apologies for the delay; I am currently filling out a lot of job applications. From the profile run itself:

   Ordered by: internal time
   List reduced from 3648 to 30 due to restriction <30>

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
2089084/15 25049.661    0.012    0.000    0.000 {method 'acquire' of '_thread.lock' objects}
  30084/0 1137.027    0.038    0.000          {method 'recv' of '_socket.socket' objects}
     7980  312.489    0.039  312.489    0.039 /Users/petergreis/Library/CloudStorage/Dropbox/Jamison/project_2/mlx/utils_mlx.py:84(compute_iou)
    35220   38.974    0.001  141.884    0.004 /Users/petergreis/Library/CloudStorage/Dropbox/Jamison/project_2/mlx/eth_mugs_dataset_mlx.py:52(__getitem__)
  1656865   25.759    0.000   25.759    0.000 {method 'read' of '_io.BufferedReader' objects}
    70560   18.426    0.000   18.426    0.000 {method 'decode' of 'ImagingDecoder' objects}
    71130   14.731    0.000   14.738    0.000 {built-in method _io.open}
    70560    9.675    0.000    9.675    0.000 {method 'astype' of 'numpy.ndarray' objects}
485334/46    5.498    0.000 26710.962  580.673 /Users/petergreis/anaconda3/envs/unet/lib/python3.12/threading.py:323(wait)
    11520    4.687    0.000  141.198    0.012 /Users/petergreis/anaconda3/envs/unet/lib/python3.12/site-packages/mlxim/data/loader.py:63(__iter__)
495723/309    4.289    0.000 29385.671   95.099 /Users/petergreis/anaconda3/envs/unet/lib/python3.12/queue.py:154(get)
   246660    4.152    0.000    4.152    0.000 {method 'encode' of 'ImagingEncoder' objects}
    30261    4.014    0.000 39006.498    1.289 /Users/petergreis/anaconda3/envs/unet/lib/python3.12/site-packages/wandb/sdk/interface/router_sock.py:25(_read_message)
35019/34890    2.941    0.000    4.112    0.000 {built-in method builtins.sum}
    70763    2.811    0.000    2.811    0.000 {built-in method posix.getcwd}
      120    2.715    0.023    2.715    0.023 {built-in method numpy.array}
    87105    2.691    0.000    6.215    0.000 /Users/petergreis/.vscode/extensions/ms-python.debugpy-2024.6.0-darwin-arm64/bundled/libs/debugpy/_vendored/pydevd/pydevd.py:1681(process_internal_commands)
   143694    2.546    0.000    4.610    0.000 /Users/petergreis/anaconda3/envs/unet/lib/python3.12/site-packages/mlx/nn/layers/normalization.py:337(__call__)
    70440    2.211    0.000    8.545    0.000 /Users/petergreis/anaconda3/envs/unet/lib/python3.12/site-packages/PIL/Image.py:686(__array_interface__)
   141000    2.086    0.000   32.431    0.000 /Users/petergreis/anaconda3/envs/unet/lib/python3.12/site-packages/PIL/ImageFile.py:175(load)
    66104    2.062    0.000    2.062    0.000 {method 'write' of '_io.BufferedRandom' objects}
   284523    1.913    0.000    1.913    0.000 {built-in method posix.lstat}
   143694    1.819    0.000    1.819    0.000 /Users/petergreis/anaconda3/envs/unet/lib/python3.12/site-packages/mlx/nn/layers/normalization.py:319(_calc_stats)
    35220    1.464    0.000    4.182    0.000 /Users/petergreis/anaconda3/envs/unet/lib/python3.12/site-packages/PIL/JpegImagePlugin.py:360(_open)
187721/19    1.347    0.000    0.293    0.015 /Users/petergreis/anaconda3/envs/unet/lib/python3.12/threading.py:637(wait)
    70560    1.199    0.000    9.405    0.000 /Users/petergreis/anaconda3/envs/unet/lib/python3.12/site-packages/PIL/Image.py:3292(_open_core)
   151677    1.164    0.000    1.397    0.000 /Users/petergreis/anaconda3/envs/unet/lib/python3.12/site-packages/mlx/nn/layers/convolution.py:132(__call__)
    70560    1.122    0.000    1.122    0.000 {method 'copy' of 'ImagingCore' objects}
    37827    1.059    0.000    1.059    0.000 {built-in method posix.stat}
    11662    1.057    0.000    4.542    0.000 /Users/petergreis/anaconda3/envs/unet/lib/python3.12/site-packages/wandb/sdk/wandb_run.py:1487(_console_raw_callback)


Execution time: 26715.3401 seconds

The model itself looks like this:

Net_model_mlx(
  (inc): DoubleConv(
    (double_conv): Sequential(
      (layers.0): Conv2d(3, 64, kernel_size=(3,), stride=(1, 1), padding=(1, 1), dilation=1, bias=True)
      (layers.1): BatchNorm(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (layers.2): ReLU()
      (layers.3): Conv2d(64, 64, kernel_size=(3,), stride=(1, 1), padding=(1, 1), dilation=1, bias=True)
      (layers.4): BatchNorm(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (layers.5): ReLU()
    )
  )
  (encode1): Encode(
    (maxpool_conv): Sequential(
      (layers.0): MaxPool2d(kernel_size=(2, 2), stride=(2, 2), padding=(0, 0))
      (layers.1): DoubleConv(
        (double_conv): Sequential(
          (layers.0): Conv2d(64, 128, kernel_size=(3,), stride=(1, 1), padding=(1, 1), dilation=1, bias=True)
          (layers.1): BatchNorm(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          (layers.2): ReLU()
          (layers.3): Conv2d(128, 128, kernel_size=(3,), stride=(1, 1), padding=(1, 1), dilation=1, bias=True)
          (layers.4): BatchNorm(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          (layers.5): ReLU()
        )
      )
    )
  )
  (encode2): Encode(
    (maxpool_conv): Sequential(
      (layers.0): MaxPool2d(kernel_size=(2, 2), stride=(2, 2), padding=(0, 0))
      (layers.1): DoubleConv(
        (double_conv): Sequential(
          (layers.0): Conv2d(128, 256, kernel_size=(3,), stride=(1, 1), padding=(1, 1), dilation=1, bias=True)
          (layers.1): BatchNorm(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          (layers.2): ReLU()
          (layers.3): Conv2d(256, 256, kernel_size=(3,), stride=(1, 1), padding=(1, 1), dilation=1, bias=True)
          (layers.4): BatchNorm(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          (layers.5): ReLU()
        )
      )
    )
  )
  (encode3): Encode(
    (maxpool_conv): Sequential(
      (layers.0): MaxPool2d(kernel_size=(2, 2), stride=(2, 2), padding=(0, 0))
      (layers.1): DoubleConv(
        (double_conv): Sequential(
          (layers.0): Conv2d(256, 512, kernel_size=(3,), stride=(1, 1), padding=(1, 1), dilation=1, bias=True)
          (layers.1): BatchNorm(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          (layers.2): ReLU()
          (layers.3): Conv2d(512, 512, kernel_size=(3,), stride=(1, 1), padding=(1, 1), dilation=1, bias=True)
          (layers.4): BatchNorm(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          (layers.5): ReLU()
        )
      )
    )
  )
  (encode4): Encode(
    (maxpool_conv): Sequential(
      (layers.0): MaxPool2d(kernel_size=(2, 2), stride=(2, 2), padding=(0, 0))
      (layers.1): DoubleConv(
        (double_conv): Sequential(
          (layers.0): Conv2d(512, 1024, kernel_size=(3,), stride=(1, 1), padding=(1, 1), dilation=1, bias=True)
          (layers.1): BatchNorm(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          (layers.2): ReLU()
          (layers.3): Conv2d(1024, 1024, kernel_size=(3,), stride=(1, 1), padding=(1, 1), dilation=1, bias=True)
          (layers.4): BatchNorm(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          (layers.5): ReLU()
        )
      )
    )
  )
  (decode1): Decode(
    (decode): ConvTranspose2d(1024, 512, kernel_size=(2,), stride=(2, 2), padding=(0, 0), dilation=1, bias=True)
    (conv): DoubleConv(
      (double_conv): Sequential(
        (layers.0): Conv2d(1024, 512, kernel_size=(3,), stride=(1, 1), padding=(1, 1), dilation=1, bias=True)
        (layers.1): BatchNorm(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (layers.2): ReLU()
        (layers.3): Conv2d(512, 512, kernel_size=(3,), stride=(1, 1), padding=(1, 1), dilation=1, bias=True)
        (layers.4): BatchNorm(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (layers.5): ReLU()
      )
    )
  )
  (decode2): Decode(
    (decode): ConvTranspose2d(512, 256, kernel_size=(2,), stride=(2, 2), padding=(0, 0), dilation=1, bias=True)
    (conv): DoubleConv(
      (double_conv): Sequential(
        (layers.0): Conv2d(512, 256, kernel_size=(3,), stride=(1, 1), padding=(1, 1), dilation=1, bias=True)
        (layers.1): BatchNorm(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (layers.2): ReLU()
        (layers.3): Conv2d(256, 256, kernel_size=(3,), stride=(1, 1), padding=(1, 1), dilation=1, bias=True)
        (layers.4): BatchNorm(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (layers.5): ReLU()
      )
    )
  )
  (decode3): Decode(
    (decode): ConvTranspose2d(256, 128, kernel_size=(2,), stride=(2, 2), padding=(0, 0), dilation=1, bias=True)
    (conv): DoubleConv(
      (double_conv): Sequential(
        (layers.0): Conv2d(256, 128, kernel_size=(3,), stride=(1, 1), padding=(1, 1), dilation=1, bias=True)
        (layers.1): BatchNorm(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (layers.2): ReLU()
        (layers.3): Conv2d(128, 128, kernel_size=(3,), stride=(1, 1), padding=(1, 1), dilation=1, bias=True)
        (layers.4): BatchNorm(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (layers.5): ReLU()
      )
    )
  )
  (decode4): Decode(
    (decode): ConvTranspose2d(128, 64, kernel_size=(2,), stride=(2, 2), padding=(0, 0), dilation=1, bias=True)
    (conv): DoubleConv(
      (double_conv): Sequential(
        (layers.0): Conv2d(128, 64, kernel_size=(3,), stride=(1, 1), padding=(1, 1), dilation=1, bias=True)
        (layers.1): BatchNorm(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (layers.2): ReLU()
        (layers.3): Conv2d(64, 64, kernel_size=(3,), stride=(1, 1), padding=(1, 1), dilation=1, bias=True)
        (layers.4): BatchNorm(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (layers.5): ReLU()
      )
    )
  )
  (outc): OutConv(
    (conv): Conv2d(64, 1, kernel_size=(1,), stride=(1, 1), padding=(0, 0), dilation=1, bias=True)
  )
  (sigmoid): Sigmoid()
)

Profiling itself is pretty straightforward, implemented so:

    start_time = time.time()
    # train(ckpt_dir, train_data_root, val_data_root)

    # Profile the training function
    profile = cProfile.Profile()
    profile.enable()
    train(ckpt_dir, train_data_root, val_data_root)
    profile.disable()

    # Sort and print profiling results
    stats = pstats.Stats(profile).sort_stats('tottime')
    stats.print_stats(30)  # Print top 30 results

    end_time = time.time()

    execution_time = end_time - start_time
    print(f"Execution time: {execution_time:.4f} seconds")