GPU使用情况

作者: Jacob Schreiber 联系方式: jmschreiber91@gmail.com

由于pomegranate模型都是torch.nn.Module的实例，您可以像操作其他PyTorch模型一样对其进行任何操作。这包括使用GPU或PyTorch支持的任何其他设备，方法调用完全相同。在这里，我们将了解如何使用GPU来加速训练和推理。

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

概述

pomegranate中的所有模型和方法都可以像其他PyTorch模型一样使用GPU加速。让我们看看实际效果。

from pomegranate.distributions import Normal

X = torch.randn(5000, 5)
dist = Normal().fit(X)

dist.means, dist.covs

(Parameter containing:
 tensor([-0.0236, -0.0205, -0.0137,  0.0102,  0.0062]),
 Parameter containing:
 tensor([[ 1.0031e+00, -7.6397e-04,  1.2268e-03,  1.8175e-02, -1.2175e-02],
         [-7.6397e-04,  1.0070e+00,  1.0800e-02,  1.4855e-02,  2.5852e-02],
         [ 1.2268e-03,  1.0800e-02,  1.0058e+00,  7.3037e-03, -1.7321e-03],
         [ 1.8175e-02,  1.4855e-02,  7.3037e-03,  1.0094e+00, -7.5648e-03],
         [-1.2175e-02,  2.5852e-02, -1.7321e-03, -7.5648e-03,  9.8117e-01]]))

我们只需要对数据和模型使用.cuda()方法或.to(device)方法即可。与PyTorch类似，模型不会自动将数据移动到GPU上，您需要自行完成这个操作。

dist = Normal().cuda().fit(X.cuda())

dist.means, dist.covs

(Parameter containing:
 tensor([-0.0236, -0.0205, -0.0137,  0.0102,  0.0062], device='cuda:0'),
 Parameter containing:
 tensor([[ 1.0031e+00, -7.6398e-04,  1.2268e-03,  1.8175e-02, -1.2175e-02],
         [-7.6398e-04,  1.0070e+00,  1.0800e-02,  1.4855e-02,  2.5852e-02],
         [ 1.2268e-03,  1.0800e-02,  1.0058e+00,  7.3037e-03, -1.7321e-03],
         [ 1.8175e-02,  1.4855e-02,  7.3037e-03,  1.0094e+00, -7.5648e-03],
         [-1.2175e-02,  2.5852e-02, -1.7321e-03, -7.5648e-03,  9.8117e-01]],
        device='cuda:0'))

所有模型都以相同的方式运行。

from pomegranate.gmm import GeneralMixtureModel

model = GeneralMixtureModel([Normal(), Normal()], max_iter=5, verbose=True).cuda()
model.fit(X.cuda())

[1] Improvement: 134.63671875, Time: 0.001251s
[2] Improvement: 38.76953125, Time: 0.001229s
[3] Improvement: 17.02734375, Time: 0.001173s
[4] Improvement: 9.13671875, Time: 0.001179s

GeneralMixtureModel(
  (distributions): ModuleList(
    (0-1): 2 x Normal()
  )
)

计时示例

当工作负载较为复杂时，使用GPU能带来最大效益。因此，在小型简单的工作负载（如基础概率分布）上，我们只会看到微小的性能提升。本次评估中，我们将分别测试三种场景的耗时：仅使用CPU、使用GPU（含数据传输时间）、以及数据预加载后的GPU运算。所有测试均在A100显卡上完成。

n, d = 100, 5

X = torch.randn(n, d)
X_cuda = X.cuda()

mu = torch.randn(d)
cov = torch.exp(torch.randn(d))

d = Normal(mu, cov, covariance_type='diag')
d_cuda = Normal(mu, cov, covariance_type='diag').cuda()

%timeit d.log_probability(X)
%timeit d.cuda().log_probability(X.cuda())
%timeit d_cuda.log_probability(X_cuda)

28.4 µs ± 107 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
91.3 µs ± 361 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
54.5 µs ± 170 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

对于这个极其小的工作负载，使用GPU比直接在CPU上处理要慢得多。让我们尝试一个更大的例子。

n, d = 10000, 50

X = torch.randn(n, d)
X_cuda = X.cuda()

mu = torch.randn(d)
cov = torch.exp(torch.randn(d))

d = Normal(mu, cov, covariance_type='diag')
d_cuda = Normal(mu, cov, covariance_type='diag').cuda()

%timeit d.log_probability(X)
%timeit d.cuda().log_probability(X.cuda())
%timeit d_cuda.log_probability(X_cuda)

131 µs ± 7.32 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
256 µs ± 108 ns per loop (mean ± std. dev. of 7 runs, 1000 loops each)
54.2 µs ± 68.3 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

看起来CPU时间增加了好几倍，但GPU时间变化不大。这是因为使用GPU有相对较高的固定成本，但随着数据规模增大，可变成本的上升速度远没有那么快。让我们尝试一个更大的例子。

n, d = 100000, 5000

X = torch.randn(n, d)
X_cuda = X.cuda()

mu = torch.randn(d)
cov = torch.exp(torch.randn(d))

d = Normal(mu, cov, covariance_type='diag')
d_cuda = Normal(mu, cov, covariance_type='diag').cuda()

%timeit d.log_probability(X)
%timeit d.cuda().log_probability(X.cuda())
%timeit d_cuda.log_probability(X_cuda)

844 ms ± 90 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
213 ms ± 639 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)
11.7 ms ± 2.97 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

我们可以看到预期的结果：尽管需要将数据传输到GPU并传回，但对于大型数据集来说，这比使用CPU更快。如果您处于所有数据都已加载到GPU的环境中，还能获得巨大的速度提升。

现在，如果您有一个更复杂的模型，您可以解锁更大的速度提升。

from pomegranate.kmeans import KMeans

model1 = KMeans(512)
model2 = KMeans(512)

%timeit -n 1 -r 1 model1.fit(X)
%timeit -n 1 -r 1 model2.cuda().fit(X.cuda())

10.1 s ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)
643 ms ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)

del X, model1, model2

看起来速度明显更快。

现在，让我们尝试一个更复杂的模型：密集隐马尔可夫模型。

from pomegranate.hmm import DenseHMM

n, l, d = 1000, 25, 15
X = torch.randn(n, l, d)

k = 256

dists1, dists2 = [], []
for i in range(k):
    mu = torch.randn(d)
    covs = torch.exp(torch.randn(d))

    dist1 = Normal(mu, covs, covariance_type='diag')
    dist2 = Normal(mu, covs, covariance_type='diag').cuda()

    dists1.append(dist1)
    dists2.append(dist2)


model1 = DenseHMM(dists1, max_iter=3)
model2 = DenseHMM(dists2, max_iter=3).cuda()

X_cuda = X.cuda()

%timeit -n 1 -r 1 model1.fit(X)
%timeit -n 1 -r 1 model2.fit(X_cuda)

8.23 s ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)
1.08 s ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)