PyTorch非阻塞传输参数详解

torch.cuda.Event() 是PyTorch中用于精确测量GPU操作时间和实现细粒度同步的工具。

基本用法

创建CUDA Event

import torch

# 创建CUDA事件
start_event = torch.cuda.Event(enable_timing=True)  # 启用计时功能
end_event = torch.cuda.Event(enable_timing=True)

record() 方法

record() 方法在GPU的当前流中插入一个事件记录点。

基本record用法

# 在GPU流中记录事件
start_event.record()  # 记录开始时间点

# 执行一些GPU操作
x = torch.randn(1000, 1000, device='cuda')
y = x * x + 2 * x + 1

end_event.record()  # 记录结束时间点

synchronize() 方法

synchronize() 方法阻塞CPU执行，直到事件完成。

基本synchronize用法

# 等待事件完成
end_event.synchronize()  # CPU等待，直到end_event被GPU触发

# 现在可以安全地读取计时结果
elapsed_time = start_event.elapsed_time(end_event)  # 毫秒
print(f"GPU操作耗时: {elapsed_time:.2f} ms")

实际应用示例

1. 测量GPU操作时间

def measure_gpu_operation():
    # 创建计时事件
    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)
    
    # 创建测试数据
    a = torch.randn(10000, 10000, device='cuda')
    b = torch.randn(10000, 10000, device='cuda')
    
    # 记录开始并执行矩阵乘法
    start.record()
    c = torch.mm(a, b)  # 矩阵乘法
    end.record()
    
    # 等待计算完成
    torch.cuda.synchronize()  # 或者使用 end.synchronize()
    
    # 计算耗时
    elapsed = start.elapsed_time(end)
    print(f"矩阵乘法耗时: {elapsed:.2f} ms")
    
    return elapsed

2. 测量数据传输时间

def measure_data_transfer():
    start_event = torch.cuda.Event(enable_timing=True)
    end_event = torch.cuda.Event(enable_timing=True)
    
    # 创建CPU张量（固定内存以获得最佳传输性能）
    cpu_tensor = torch.randn(5000, 5000).pin_memory()
    
    # 测量数据传输时间
    start_event.record()
    gpu_tensor = cpu_tensor.to('cuda', non_blocking=True)
    end_event.record()
    
    # 等待传输完成
    end_event.synchronize()
    
    elapsed = start_event.elapsed_time(end_event)
    print(f"CPU->GPU数据传输耗时: {elapsed:.2f} ms")
    
    return elapsed

3. 训练循环中的精确计时

def training_loop_with_timing(dataloader, model, optimizer):
    model = model.cuda()
    
    # 创建事件用于计时
    data_transfer_start = torch.cuda.Event(enable_timing=True)
    data_transfer_end = torch.cuda.Event(enable_timing=True)
    forward_start = torch.cuda.Event(enable_timing=True)
    forward_end = torch.cuda.Event(enable_timing=True)
    backward_start = torch.cuda.Event(enable_timing=True)
    backward_end = torch.cuda.Event(enable_timing=True)
    
    for batch_idx, (data, target) in enumerate(dataloader):
        # 测量数据传输时间
        data_transfer_start.record()
        data = data.to('cuda', non_blocking=True)
        target = target.to('cuda', non_blocking=True)
        data_transfer_end.record()
        
        # 测量前向传播时间
        forward_start.record()
        output = model(data)
        loss = torch.nn.functional.cross_entropy(output, target)
        forward_end.record()
        
        # 测量反向传播时间
        backward_start.record()
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        backward_end.record()
        
        # 等待所有操作完成
        torch.cuda.synchronize()
        
        # 计算各阶段耗时
        transfer_time = data_transfer_start.elapsed_time(data_transfer_end)
        forward_time = forward_start.elapsed_time(forward_end)
        backward_time = backward_start.elapsed_time(backward_end)
        
        if batch_idx % 100 == 0:
            print(f"Batch {batch_idx}: 传输={transfer_time:.2f}ms, "
                  f"前向={forward_time:.2f}ms, 反向={backward_time:.2f}ms")

4. 多流同步

def multi_stream_sync():
    # 创建多个流
    stream1 = torch.cuda.Stream()
    stream2 = torch.cuda.Stream()
    
    # 创建同步事件
    sync_event = torch.cuda.Event()
    
    with torch.cuda.stream(stream1):
        # 在流1中执行操作
        a = torch.randn(1000, 1000, device='cuda')
        b = torch.mm(a, a)  # 一些计算
        sync_event.record(stream=stream1)  # 在流1中记录事件
    
    with torch.cuda.stream(stream2):
        # 流2等待流1的事件完成
        torch.cuda.current_stream().wait_event(sync_event)
        # 现在可以安全地使用流1的结果
        c = b * 2  # 使用流1的计算结果
    
    # 等待所有操作完成
    torch.cuda.synchronize()

5. 异步操作的正确同步

def async_operations():
    # 创建事件用于同步
    data_ready = torch.cuda.Event()
    
    # 在默认流中准备数据
    data = torch.randn(1000, 1000, device='cuda')
    
    # 记录数据准备完成
    data_ready.record()
    
    # 创建新流并等待数据准备事件
    compute_stream = torch.cuda.Stream()
    with torch.cuda.stream(compute_stream):
        # 等待数据准备完成
        compute_stream.wait_event(data_ready)
        
        # 现在安全地进行计算
        result = data * data + 2
    
    # 等待计算流完成
    compute_stream.synchronize()
    
    return result

重要注意事项

1. enable_timing参数：

   # 如果需要计时，必须设置enable_timing=True
   event = torch.cuda.Event(enable_timing=True)

2. 流关联：

   # 事件与创建时的流关联，或通过record()指定流
   event.record(stream=torch.cuda.current_stream())

3. 同步选择：

   # 可以选择同步特定事件或所有GPU操作
   event.synchronize()        # 等待特定事件
   torch.cuda.synchronize()   # 等待所有GPU操作

4. 性能分析：

   # 用于性能分析的最佳实践
   def profile_operation():
       start = torch.cuda.Event(enable_timing=True)
       end = torch.cuda.Event(enable_timing=True)
       
       # 预热GPU
       warmup = torch.randn(100, 100, device='cuda')
       _ = warmup * warmup
       torch.cuda.synchronize()
       
       # 实际测量
       start.record()
       # 要测量的操作
       end.record()
       torch.cuda.synchronize()
       
       return start.elapsed_time(end)

使用torch.cuda.Event可以精确控制GPU操作的同步和计时，对于性能优化和调试非常有用。