Add AMP (automatic mixed precision) as option for training and inference

baler-collaboration / baler

Repository of Baler, a machine learning based data compression tool

https://github.com/baler-collaboration/baler.github.io

Apache License 2.0

33 stars 31 forks source link

Add AMP (automatic mixed precision) as option for training and inference #304

Open neogyk opened 1 year ago

neogyk commented 1 year ago

Usage of half-precision speed ups the training process of DNN, as it allocates less memory.

neogyk commented 1 year ago

Currently, the baler can be executed only with full precision numbers. The mixed precision is not working due to the error of mismatch of the data types during the matrix multiplication in the forward pass.