Mixed Precision and Other Data Types¶
author: Jacob Schreiber contact: email@example.com
Because pomegranate models are all instances of
torch.nn.Module, you can do anything with them that you could do with other PyTorch models. In the first tutorial, we saw how this means that one can use GPUs in exactly the same way that one would with their other PyTorch models. However, this also means that all the great things built-in for during half precision, quantization, automatic mixed precision (AMP), etc., can also be used in pomegranate.
import os os.environ['CUDA_VISIBLE_DEVICES'] = '0' %pylab inline import seaborn; seaborn.set_style('whitegrid') import torch numpy.random.seed(0) numpy.set_printoptions(suppress=True) %load_ext watermark %watermark -m -n -p numpy,scipy,torch,pomegranate
Populating the interactive namespace from numpy and matplotlib numpy : 1.23.4 scipy : 1.9.3 torch : 1.13.0 pomegranate: 1.0.0 Compiler : GCC 11.2.0 OS : Linux Release : 4.15.0-208-generic Machine : x86_64 Processor : x86_64 CPU cores : 8 Architecture: 64bit
float16 (half) Precision¶
Doing operations at half precision is just the same as it would by in PyTorch: you can use the
.half() method or the
.to(torch.float16) method. However, very few operations seem to be supported for half precision, including the log and sqrt methods not being supported for some reason? So, until more operations are supported, you will probably be using other methods.
More operations seem to be supported for
from pomegranate.distributions import Normal X = torch.randn(1000, 5) d = Normal(covariance_type='diag').fit(X) d.means, d.covs
(Parameter containing: tensor([-0.0076, 0.0164, 0.0164, -0.0033, 0.0105]), Parameter containing: tensor([1.0291, 1.0491, 0.9661, 0.8947, 1.0778]))
X = X.to(torch.bfloat16) d = Normal(covariance_type='diag').to(torch.bfloat16).fit(X) d.means, d.covs
(Parameter containing: tensor([-0.0075, 0.0165, 0.0165, -0.0034, 0.0103], dtype=torch.bfloat16), Parameter containing: tensor([1.0312, 1.0469, 0.9688, 0.8945, 1.0781], dtype=torch.bfloat16))
However, not all operations are supported for
torch.bfloat16 either, including the cholesky decomposition used in full covariance normal distributions.
Although not all operations support all data types, all models and methods (inference and training) support them to the extent that the underlying operations allow. For instance, we can just as easily use a mixture model with
bfloat16 data types as with full floats.
from pomegranate.gmm import GeneralMixtureModel model = GeneralMixtureModel([Normal(covariance_type='diag'), Normal(covariance_type='diag')], verbose=True) model = model.to(torch.bfloat16) model.fit(X)
 Improvement: 32.0, Time: 0.001619s  Improvement: 32.0, Time: 0.001082s  Improvement: 0.0, Time: 0.00106s
GeneralMixtureModel( (distributions): ModuleList( (0-1): 2 x Normal() ) )
And we can use the resulting trained model to make predictions at whatever resolution we’d like.
y_hat = model.predict_proba(X) y_hat, y_hat.dtype
(tensor([[0.8281, 0.1680], [0.2451, 0.7539], [0.0874, 0.9375], ..., [0.6445, 0.3574], [0.3145, 0.6875], [0.2695, 0.7305]], dtype=torch.bfloat16), torch.bfloat16)
model = model.to(torch.float32) y_hat = model.predict_proba(X) y_hat, y_hat.dtype
(tensor([[0.8299, 0.1701], [0.2463, 0.7537], [0.0793, 0.9207], ..., [0.6496, 0.3504], [0.3148, 0.6852], [0.2788, 0.7212]]), torch.float32)
Automatic Mixed Precision¶
An automatic way to get around some of these issues is to use AMP so that operations which can work at lower precision are cast and others are not. Keeping up with the theme, doing this is exactly the same as using AMP with your other PyTorch models.
X = torch.randn(1000, 50).cuda() model = GeneralMixtureModel([Normal(), Normal()]).cuda() with torch.autocast('cuda', dtype=torch.float16): model.fit(X)
This would have crashed if you tried to run
model.fit alone because of the unsupported Cholesky decomposition.
Unfortunately, because pomegranate uses a wide range of operations to implement the underlying models, the speedups from using mixed precision are inconsistent. It may be worth trying out in your application but the speedups observed in training neural networks are not guaranteed here because not all operations are supported – and if they are supported, they may not be optimized to be faster. Basically, because AMP will fall back on normal precision for many operations, the entire method may end up not being significantly faster in practice.
X = torch.randn(10000, 500).cuda() model = GeneralMixtureModel([Normal(covariance_type='diag') for i in range(10)], max_iter=5, verbose=True).cuda() with torch.autocast('cuda', dtype=torch.bfloat16): model.fit(X) print() model = GeneralMixtureModel([Normal(covariance_type='diag') for i in range(10)], max_iter=5, verbose=True).cuda() model.fit(X)
 Improvement: 1015.0, Time: 0.1042s  Improvement: 417.5, Time: 0.1044s  Improvement: 310.0, Time: 0.1041s  Improvement: 243.5, Time: 0.1041s  Improvement: 1232.5, Time: 0.1039s  Improvement: 613.5, Time: 0.1042s  Improvement: 446.5, Time: 0.1043s  Improvement: 344.0, Time: 0.1043s
GeneralMixtureModel( (distributions): ModuleList( (0-9): 10 x Normal() ) )