利用Nisght System对简单的深度学习训练神经网络进行性能分析

Nsight System常用的使用方式时，服务器上生成报告文件，然后在自己的笔记本上查看。因此，服务器和PC上都需要下载（下载的版本号必须相同，否则PC上无法打开在服务器上生成的文件）。下载链接：https://developer.nvidia.com/nsight-systems/get-started

直接在mnist的训练文件上增加nvtx的注释：

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

import nvtx


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.Dropout(0.25)
        self.dropout2 = nn.Dropout(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.relu(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()

    nvtx.push_range("Data loading")
    for batch_idx, (data, target) in enumerate(train_loader):
        # 此时 Data loading 结束
        nvtx.pop_range()

        with nvtx.annotate(f"Batch {batch_idx}", color="green"):

            with nvtx.annotate("Copy to device", color="yellow"):
                data, target = data.to(device), target.to(device)

            with nvtx.annotate("Forward pass", color="blue"):
                optimizer.zero_grad()
                output = model(data)
                loss = F.nll_loss(output, target)

            with nvtx.annotate("Backward pass", color="red"):
                loss.backward()
                optimizer.step()

        nvtx.push_range("Data loading")

    nvtx.pop_range()


def test(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.0 * 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-accel", action="store_true", help="disables accelerator")
    parser.add_argument(
        "--dry-run", action="store_true", help="quickly check a single pass"
    )
    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", help="For Saving the current Model"
    )
    args = parser.parse_args()

    use_accel = not args.no_accel and torch.accelerator.is_available()

    torch.manual_seed(args.seed)

    if use_accel:
        device = torch.accelerator.current_accelerator()
    else:
        device = torch.device("cpu")

    train_kwargs = {"batch_size": args.batch_size}
    test_kwargs = {"batch_size": args.test_batch_size}
    if use_accel:
        accel_kwargs = {
            "num_workers": 1,
            "persistent_workers": True,
            "pin_memory": True,
            "shuffle": True,
        }
        train_kwargs.update(accel_kwargs)
        test_kwargs.update(accel_kwargs)

    transform = transforms.Compose(
        [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]
    )
    dataset1 = datasets.MNIST("../data", train=True, download=True, transform=transform)
    dataset2 = datasets.MNIST("../data", train=False, transform=transform)
    train_loader = torch.utils.data.DataLoader(dataset1, **train_kwargs)
    test_loader = torch.utils.data.DataLoader(dataset2, **test_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):
        nvtx.push_range(f"Epoch {epoch}")

        with nvtx.annotate("Training", color="dodgerblue"):
            train(args, model, device, train_loader, optimizer, epoch)

        with nvtx.annotate("Testing", color="dodgerblue"):
            test(model, device, test_loader)

        scheduler.step()
        nvtx.pop_range()
    if args.save_model:
        torch.save(model.state_dict(), "mnist_cnn.pt")


if __name__ == "__main__":
    main()

使用不同的batch size对比差异。run.sh：

# BATCH_SIZE=32
BATCH_SIZE=64
rm -rf baseline-${BATCH_SIZE}.*
nsys profile -t cuda,osrt,nvtx -o baseline-${BATCH_SIZE} -w true python main.py --epochs 1 --batch-size ${BATCH_SIZE}

然后执行脚本文件 bash run.sh
得到两个文件：baseline-64.nsys-rep baseline-32.nsys-rep。下载到本地，查看，可以发现，GPU空闲的时间在于Dataloader加载数据，因此降低batch size，可以提高GPU使用率：

• baseline-64.nsys-rep：
  

• baseline-32.nsys-rep