Nsys/TorchProfiler/NCU 使用

# 导入所有需要的库
import torch
from torch import nn
from torch.utils.data import DataLoader, random_split
from torchvision import datasets
# 导入 transforms 模块，用于数据增强
from torchvision import transforms 
import time
import nvtx

# --- 全局配置参数 ---

# 设置计算设备
device = 「cuda」 if torch.cuda.is_available() else 「cpu」

# 设置训练的总轮数
num_epoch = 1

# 设置批次大小
batch_size = 128

# 设置数据加载器工作进程数量
# 切换这里的数值 (0, 2, 4 等) 并用 nsys 观察差异
num_workers = 32 


# --- 定义神经网络模型 ---

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Conv2d(1, 32, 5), nn.ReLU(),
            nn.Conv2d(32, 32, 5), nn.ReLU(),
            nn.BatchNorm2d(32),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 3), nn.ReLU(),
            nn.Conv2d(64, 64, 3), nn.ReLU(),
            nn.BatchNorm2d(64),
            nn.MaxPool2d(2),
            nn.Flatten(1),
            nn.Linear(576, 10)
        )

    @nvtx.annotate(「forward」, color=「blue」)
    def forward(self, x):
        return self.layers(x)

# --- 定义训练和测试的单步函数 ---

def train_step(dataloader, model, loss_fn, opt):
    model.train()
    for batch, (X, y) in enumerate(dataloader):
        X, y = X.to(device), y.to(device)
        
        # 1. 前向传播
        pred = model(X)
        loss = loss_fn(pred, y)
        
        # 2. 反向传播和优化
        with nvtx.annotate(「backward」, color=「red」):
            loss.backward()
        
        with nvtx.annotate(「opt」, color=「green」):
            opt.step()
            opt.zero_grad()

def test_step(dataloader, model, loss_fn, opt):
    size = len(dataloader.dataset)
    num_batches = len(dataloader)
    model.eval()
    test_loss, correct = 0.0, 0.0
    
    with torch.no_grad():
        for X, y in dataloader:
            X, y = X.to(device), y.to(device)
            pred = model(X)
            test_loss += loss_fn(pred, y)
            correct += (pred.argmax(1) == y).type(torch.float).sum()
    
    test_loss /= num_batches
    correct /= size
    print(f『Accuracy: {(100*correct):.1f}%, AVG loss: {test_loss:.8f} \n』)

# --- 制造 CPU 瓶颈的自定义 Transform ---
class SlowTransform:
    「」「一个用于演示的慢变换，模拟复杂的 CPU 预处理」「」
    def __call__(self, img):
        # 模拟一个复杂的、CPU 密集型的图像处理过程
        # 注意：在真实项目中不要使用 time.sleep
        time.sleep(0.001) 
        return img

# --- 主执行流程函数 ---

def Pipe():
    # --- 1. 数据准备 ---
    
    # 核心：定义一个 CPU 密集的转换组合 (Data Augmentation)
    heavy_transforms = transforms.Compose([
        transforms.RandomRotation(15),       # 随机旋转 (CPU 密集)
        transforms.RandomCrop(28, padding=4), # 随机裁剪 (CPU 密集)
        transforms.ToTensor(),              # 转换为张量
        SlowTransform()                     # 自定义的慢变换
    ])

    # 下载并加载训练数据集，应用 CPU 密集型变换
    training_data = datasets.MNIST(
        root=「data」, train=True, download=True, transform=heavy_transforms,
    )

    # 下载并加载测试数据集
    test_data = datasets.MNIST(
        root=「data」, train=False, download=True, transform=transforms.ToTensor(),
    )

    total_samples = len(training_data)
    # 使用一半数据进行训练
    training_samples, _ = random_split(training_data, [total_samples//2, total_samples-total_samples//2])
    
    # 创建训练数据加载器 (开启 pin_memory 可以加速数据从 CPU 到 GPU 的传输)
    train_dataloader = DataLoader(
        training_samples, batch_size=batch_size, num_workers=num_workers, pin_memory=True
    )

    # 创建测试数据加载器
    test_dataloader = DataLoader(
        test_data, batch_size=batch_size, num_workers=num_workers, pin_memory=True
    )

    # 打印数据形状
    for X, y in test_dataloader:
        print(f「Shape of X [N, C, H, W]: {X.shape}」) 
        print(f「Shape of y: {y.shape} {y.dtype}」) 
        break 

    # --- 2. 模型、损失函数和优化器初始化 ---
    model = Model().to(device)
    loss_fn = nn.CrossEntropyLoss()
    opt = torch.optim.Adam(model.parameters(), lr=1e-4)

    # --- 3. 训练循环 ---

    # 预热阶段：处理前几个批次以消除启动开销，这对性能分析很重要
    print(「Starting warmup...」)
    for _ in range(5): 
        # 确保数据加载器可以提供数据
        try:
            X, y = next(iter(train_dataloader))
        except StopIteration:
            break

        X, y = X.to(device), y.to(device)
        opt.zero_grad()
        loss = loss_fn(model(X), y)
        loss.backward()
        opt.step()
    
    print(「Warmup finished. Starting profiling...」)
    torch.cuda.synchronize() # 确保预热结束

    for epoch in range(num_epoch):
        print(f「Epoch {epoch+1}\n-------------------------------」)
        now = time.time()
        
        # 训练步骤
        train_step(dataloader=train_dataloader, model=model, loss_fn=loss_fn, opt=opt)
        
        # 确保 GPU 所有任务完成
        torch.cuda.synchronize() 
        epoch_duration = time.time() - now
        
        print(f「Epoch duration: {epoch_duration:.2f} s」)
        print(f「avg batch train speed: {(len(train_dataloader)/epoch_duration):.2f} it/s」)
        
        now = time.time()
        test_step(dataloader=test_dataloader, model=model, loss_fn=loss_fn, opt=opt)
        torch.cuda.synchronize()
        test_duration = time.time() - now
        print(f「teststep speed: {(len(test_dataloader)/test_duration):.2f} it/s」)

if __name__ == 「__main__」:
    Pipe()