Speed up inference or/and training of a 3D deep neural network (U-net) for a regression task

Hi all,

I am working on a (volumetric) regression task using a 3D deep neural network.

Its architecture is based on the 3D U-net model provided by the output of Matlab's function unet3dLayers(). I modified its architecture by switching the upscaling transposedConv3dLayer layers with resize3dLayer layers. Furthermore, I removed the PixelClassificationLayer and defined a custom deep learning loop and loss function as described here: https://de.mathworks.com/help/deeplearning/ug/train-network-using-custom-training-loop.html.

The training and inference work well, but their duration is too long for the later use case. Thus, I tried to use the "Deep Network Quantizer" to speed up the inference time, but the toolbox does not support 3D layers. Also, other optimisation strategies for inference/training do not seem to be supported for 3D layers.

So my questions is: Is there any other technique to speed up the inference and/or the training of this type of network architecture?

Many thanks in advance!

Answer 1

[Joss Knight]

Have you tried using dlaccelerate? As well as ensuring any Custom Layers are using the Acceleratable mixin?

[Emanuel Maus]

Thank you for your answer and the great advice.

I managed to speed up foward- and backpropagation (model_loss function) with the help of the dlaccelerate function and I measured a reduction of the training time by almost 10%.Is there any other part in a custom training loop (for training a dlnetwork), where I can use the dlaccelerate function?

About the Acceleratable mixin: I used only Matlab's own layers for the U-net, so I guess they naturally inherent from the nnet.layer.Acceleratable class, don't they? If not, where can I look if the layer uses it?

Furthermore, does this Acceleratable mixin exist in all the recent Matlab versions?

Thanks a lot in advance!

[Joss Knight]

Hi Emanuel. Yes, if you do not have any Custom Layers then everything is already Acceleratable. The Acceleratable mixin was introduced in R2022a.

You can accelerate most code that uses dlarrays (that is traceable), so for instance you can accelerate any preprocessing you do to the input before passing it to your modelLoss function, validation during training or testing and post-processing after training. We are working on extending acceleration to the parameter update part of training.