PythonとDeep learningとそして少しRustを...。
2020/02/05追記
下の記事は、Pythonを使うとGPU使用率が低くなってしまい、計算速度が遅いという話をしています。しかし、rustがデータ全てをgpuに移動させていることや、pytorchのtorch.backends.cudnn.benchmark=Trueを指定しいないことなどから再実験しました。今回の実験のpythonでは、あらかじめバッチを全てgpuに移動させています。バッチサイズ256で50エポック計算した時の1エポックの平均計算時間を下記の表は示しています。
| RTX2060 | |
|---|---|
| python | 5.32s(ave gpu util~80%) |
python(torch.backends.cudnn.benchmark=True) |
4.24s(ave gpu util~80%) |
| rust | 3.80s(ave gpu util~90%) |
以前の記事
「機械学習をやるならPython」これは今や切っても切り離せない関係になってしまったかもしれない。Pythonのユーザーフレンドリーな仕様はプログラム未経験者でも使い易く、覚えやすい。一方で、コードをインタープリンタで解釈し逐次実行していることから、コンパイル型言語の実行コードに比べ数百倍遅くなってしまう。なので、Pthonしか知らない人はPythonの遅さに気付いたとき驚愕するかも知れない。データサイエンティストにとってはプログラム修得に時間をかけるよりも実験的なコードをいかにたくさん試行できるかが重要であり、実行時間より開発時間の短縮を選択した結果なのかもと考えています。さらに、多くの計算速度を必要とするライブラリのバックエンドは、C++で書かれており、ユーザーが計算速度を意識する必要はもうすでにないとも言えます。だがしかし、DeepLearningにおいては私は少し違う意見を持っています。
DeepLearningにおいても多くのライブラリのバックエンドはすでにC++に置き換わっています。ましてや、cuDNNを使わない手はないし、速度を求めるならcuDNNの使用は必然であります。一見、これだけなら「DeepLearningをやるならPython」という主張になんら違和感がないように思います。所詮はバックエンドがC++なのだからと。ところが、実際に計算を回してみるとGPUユーティライゼージョンが50%~なんてことはよく目にします。これはGPUの性能を発揮できていない証拠です。どうしてか...。
CUDAプログラミングのコツは、なるべく多くのWarpを走らせた方が良い。この部分はDeepLearningでいうと以下にバッチ作成を高速化するかが重要であかを意味しています。以下に面白い実験をしたので結果を紹介します。
こちらは、pytorchによるMnistトレーニングのサンプルコード
from __future__ import print_function import argparse import torch import torch.nn as nn import torch.nn.functional as F import torch.optim as optim from torchvision import datasets, transforms from torch.optim.lr_scheduler import StepLR class Net(nn.Module): def __init__(self): super(Net, self).__init__() self.conv1 = nn.Conv2d(1, 32, 3, 1) self.conv2 = nn.Conv2d(32, 64, 3, 1) self.dropout1 = nn.Dropout2d(0.25) self.dropout2 = nn.Dropout2d(0.5) self.fc1 = nn.Linear(9216, 128) self.fc2 = nn.Linear(128, 10) def forward(self, x): x = self.conv1(x) x = F.relu(x) x = self.conv2(x) x = F.max_pool2d(x, 2) x = self.dropout1(x) x = torch.flatten(x, 1) x = self.fc1(x) x = F.relu(x) x = self.dropout2(x) x = self.fc2(x) output = F.log_softmax(x, dim=1) return output def train(args, model, device, train_loader, optimizer, epoch): model.train() for batch_idx, (data, target) in enumerate(train_loader): data, target = data.to(device), target.to(device) optimizer.zero_grad() output = model(data) loss = F.nll_loss(output, target) loss.backward() optimizer.step() if batch_idx % args.log_interval == 0: print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format( epoch, batch_idx * len(data), len(train_loader.dataset), 100. * batch_idx / len(train_loader), loss.item())) def test(args, model, device, test_loader): model.eval() test_loss = 0 correct = 0 with torch.no_grad(): for data, target in test_loader: data, target = data.to(device), target.to(device) output = model(data) test_loss += F.nll_loss(output, target, reduction='sum').item() # sum up batch loss pred = output.argmax(dim=1, keepdim=True) # get the index of the max log-probability correct += pred.eq(target.view_as(pred)).sum().item() test_loss /= len(test_loader.dataset) print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format( test_loss, correct, len(test_loader.dataset), 100. * correct / len(test_loader.dataset))) def main(): # Training settings parser = argparse.ArgumentParser(description='PyTorch MNIST Example') parser.add_argument('--batch-size', type=int, default=64, metavar='N', help='input batch size for training (default: 64)') parser.add_argument('--test-batch-size', type=int, default=1000, metavar='N', help='input batch size for testing (default: 1000)') parser.add_argument('--epochs', type=int, default=14, metavar='N', help='number of epochs to train (default: 14)') parser.add_argument('--lr', type=float, default=1.0, metavar='LR', help='learning rate (default: 1.0)') parser.add_argument('--gamma', type=float, default=0.7, metavar='M', help='Learning rate step gamma (default: 0.7)') parser.add_argument('--no-cuda', action='store_true', default=False, help='disables CUDA training') parser.add_argument('--seed', type=int, default=1, metavar='S', help='random seed (default: 1)') parser.add_argument('--log-interval', type=int, default=10, metavar='N', help='how many batches to wait before logging training status') parser.add_argument('--save-model', action='store_true', default=False, help='For Saving the current Model') args = parser.parse_args() use_cuda = not args.no_cuda and torch.cuda.is_available() torch.manual_seed(args.seed) device = torch.device("cuda" if use_cuda else "cpu") kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {} train_loader = torch.utils.data.DataLoader( datasets.MNIST('../data', train=True, download=True, transform=transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) ])), batch_size=args.batch_size, shuffle=True, **kwargs) test_loader = torch.utils.data.DataLoader( datasets.MNIST('../data', train=False, transform=transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) ])), batch_size=args.test_batch_size, shuffle=True, **kwargs) model = Net().to(device) optimizer = optim.Adadelta(model.parameters(), lr=args.lr) scheduler = StepLR(optimizer, step_size=1, gamma=args.gamma) for epoch in range(1, args.epochs + 1): train(args, model, device, train_loader, optimizer, epoch) test(args, model, device, test_loader) scheduler.step() if args.save_model: torch.save(model.state_dict(), "mnist_cnn.pt") if __name__ == '__main__': main()
こちらは、pytorchlibを使ったrustのMnistの学習のサンプルコード
// CNN model. This should rearch 99.1% accuracy. use tch::{nn, nn::ModuleT, nn::OptimizerConfig, Device, Tensor}; #[derive(Debug)] struct Net { conv1: nn::Conv2D, conv2: nn::Conv2D, fc1: nn::Linear, fc2: nn::Linear, } impl Net { fn new(vs: &nn::Path) -> Net { let conv1 = nn::conv2d(vs, 1, 32, 3, Default::default()); let conv2 = nn::conv2d(vs, 32, 64, 3, Default::default()); let fc1 = nn::linear(vs, 9216, 128, Default::default()); let fc2 = nn::linear(vs, 128, 10, Default::default()); Net { conv1, conv2, fc1, fc2, } } } impl nn::ModuleT for Net { fn forward_t(&self, xs: &Tensor, train: bool) -> Tensor { xs.view([-1, 1, 28, 28]) .apply(&self.conv1) .relu() .apply(&self.conv2) .max_pool2d_default(2) .dropout_(0.5, train) .view([-1, 9216]) .apply(&self.fc1) .relu() .dropout_(0.5, train) .apply(&self.fc2) } } pub fn run() -> failure::Fallible<()> { let m = tch::vision::mnist::load_dir("data")?; let vs = nn::VarStore::new(Device::cuda_if_available()); let net = Net::new(&vs.root()); let mut opt = nn::Adam::default().build(&vs, 1e-4)?; for epoch in 1..100 { for (bimages, blabels) in m.train_iter(256).shuffle().to_device(vs.device()) { let loss = net .forward_t(&bimages, true) .cross_entropy_for_logits(&blabels); opt.backward_step(&loss); } let test_accuracy = net.batch_accuracy_for_logits(&m.test_images, &m.test_labels, vs.device(), 1024); println!("epoch: {:4} test acc: {:5.2}%", epoch, 100. * test_accuracy,); } Ok(()) }
上記のコードを私の手元にあるGPUで50epoch学習してみました。計測方法はコマンドラインツールであるtime使い計測しました。(rustはtime cargo runで実行)
| RTX2060 | RTX2080ti | |
|---|---|---|
| python | 398.25s(ave gpu util~50%) | 432.63s(ave gpu util~20%) |
| rust | 135.28s(ave gpu util~90%) | 65.03s(ave gpu util~90%) |
上記の結果は、もちろんwarpだけが原因ではないように思います。特にPythonでrtx2060よりrtx2080tiの方が計算時間が長くなっており、utilizationも下がっています。それがどうしてか...。ただ結果的にはrustの方がGPUを有効に使えています。Pythonでもバッチサイズやネットワークの深さなどを調整することでGPU utilizationをあげることも可能です。ただ、DeepLearningをやるなら速度の点からみるとPython以外も十分に選択肢にあるということ示す結果であると思います。
