PyTorchでTensorコアを使用してCNNモデルをトレーニングするための比較

Github（英語バージョン）
https://github.com/leolui2004/tensorcore

これは PyTorch フレームワークで TensorCore を使用して CNN モデルをトレーニングするための簡単な比較です。公式サイトにはサンプルスクリプトがありますが、実際のアプリケーションには多くの違いがあるはずです。

ハードウェアとソフトウェア

GPU: Tesla T4

Framework: PyTorch 1.11

nvidia Driver: 470.57.02

cuda: 11.4

パフォーマンス

テンソルコアを使用しないサンプルコード

for epoch in range(epochs):
    for input, target in zip(data, targets):
        output = net(input)
        loss = loss_fn(output, target)
        loss.backward()
        opt.step()
        opt.zero_grad()

Default precision:
Total execution time = 7.020 sec
Max memory used by tensors = 1350681600 bytes

テンソルコアを使用したサンプルコード

use_amp = True
scaler = torch.cuda.amp.GradScaler(enabled=use_amp)

for epoch in range(epochs):
    for input, target in zip(data, targets):
        with torch.cuda.amp.autocast(enabled=use_amp):
            output = net(input)
            loss = loss_fn(output, target)
        scaler.scale(loss).backward()
        scaler.step(opt)
        scaler.update()
        opt.zero_grad()

Mixed precision:
Total execution time = 2.524 sec
Max memory used by tensors = 1577232384 bytes

TensorCore を使用したものはさらに多くのメモリを使用しましたが、実行時間は TensorCore を使用しない場合の半分です。

テンソルコアを使用しない CNN を使用した MNIST（バッチサイズ 512）

trainloader = torch.utils.data.DataLoader(trainset,batch_size=512,shuffle=True,num_workers=2)

for epoch in range(20):
    running_loss = 0.0
    for i, (inputs, labels) in enumerate(trainloader, 0):
        optimizer.zero_grad()

        outputs = net(inputs.to(device))
        loss = loss_fn(outputs, labels.to(device))
        
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
    
    scheduler.step()

Total execution time = 107.228 sec
Max memory used by tensors = 417359872 bytes

テンソルコアを使用した CNN を使用したMNIST（バッチサイズ 512）

trainloader = torch.utils.data.DataLoader(trainset,batch_size=512,shuffle=True,num_workers=2)

use_amp = True
scaler = torch.cuda.amp.GradScaler(enabled=use_amp)

for epoch in range(20):
    running_loss = 0.0
    for i, (inputs, labels) in enumerate(trainloader, 0):
        optimizer.zero_grad()

        with torch.cuda.amp.autocast(enabled=use_amp):
            outputs = net(inputs.to(device))
            loss = loss_fn(outputs, labels.to(device))
        
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()

        running_loss += loss.item()
        if i % 500 == 499:
            print('[{:d}, {:5d}] loss: {:.3f}'
                    .format(epoch + 1, i + 1, running_loss / 100))
            running_loss = 0.0
    
    scheduler.step()

Total execution time = 104.948 sec
Max memory used by tensors = 175422976 bytes

テンソルコアを使用した CNN を使用した MNIST（バッチサイズ 1024）

trainloader = torch.utils.data.DataLoader(trainset,batch_size=1024,shuffle=True,num_workers=2)

use_amp = True
scaler = torch.cuda.amp.GradScaler(enabled=use_amp)

for epoch in range(20):
    running_loss = 0.0
    for i, (inputs, labels) in enumerate(trainloader, 0):
        optimizer.zero_grad()

        with torch.cuda.amp.autocast(enabled=use_amp):
            outputs = net(inputs.to(device))
            loss = loss_fn(outputs, labels.to(device))
        
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()

        running_loss += loss.item()
        if i % 500 == 499:
            print('[{:d}, {:5d}] loss: {:.3f}'
                    .format(epoch + 1, i + 1, running_loss / 100))
            running_loss = 0.0
    
    scheduler.step()

Total execution time = 100.529 sec
Max memory used by tensors = 320654848 bytes

TensorCore を使用しても、実行時間は大幅に短縮されませんでした。逆に、バッチサイズを2倍にした後でも、使用される最大メモリが元のメモリよりも少なくなり、実行時間が少し短縮されたとしても、メモリは半分に減少しました。

いくつかの理由で、TensorCore を使用したモデルの他のボトルネックはトレーニング速度を改善しませんでしたが、メモリが少ないということは、トレーニングに使用できるバッチサイズが大きいことを意味します。通常、実世界の大きなモデルのトレーニングには多くの GPU メモリー、GPU メモリーがあまりない小さな VM にこれらのモデルをデプロイする方がはるかに簡単です。