pytorch源码阅读

register_hook

参考资料：

1. [Pytorch] Tensor底层机制 https://blog.csdn.net/smartcat2010/article/details/118882020
2. cudaLauchKernel传入参数的一个提问 https://cloud.tencent.com/developer/ask/sof/107247510
3. cuda中LauchKernel的过程介绍。 https://jia.je/software/2023/10/17/clang-cuda-support/#host-代码

杂项

百度codemate生成的一个ddp例子

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.distributed import init_process_group, destroy_process_group, get_rank, get_world_size

# 假设的模型和数据集
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(10, 2)  # 示例模型

    def forward(self, x):
        return self.fc(x)

class SimpleDataset(Dataset):
    def __init__(self):
        # 初始化数据集
        self.data = torch.randn(100, 10)
        self.labels = torch.randint(0, 2, (100,))

    def __getitem__(self, index):
        return self.data[index], self.labels[index]

    def __len__(self):
        return len(self.data)

def train(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 = nn.functional.cross_entropy(output, target)
        loss.backward()
        optimizer.step()
        if batch_idx % 10 == 0:
            print(f'Train Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)} ({100. * batch_idx / len(train_loader):.0f}%)]\tLoss: {loss.item():.6f}')

def main():
    # 初始化分布式环境
    world_size = torch.cuda.device_count()
    torch.manual_seed(0)
    torch.cuda.manual_seed_all(0)
    init_process_group(backend='nccl', init_method='env://', world_size=world_size, rank=0)

    # 设置设备
    device = torch.device(f"cuda:{get_rank()}")

    # 创建模型和数据集
    model = SimpleModel().to(device)
    model = DDP(model, device_ids=[get_rank()])
    train_dataset = SimpleDataset()
    train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)

    # 优化器
    optimizer = optim.SGD(model.parameters(), lr=0.01)

    # 训练
    for epoch in range(1, 6):
        train(model, device, train_loader, optimizer, epoch)

    # 清理
    destroy_process_group()

if __name__ == '__main__':
    main()

注意：

上面的代码示例中，init_process_group 的 init_method 设置为 'env://'，这通常用于通过环境变量（如 MASTER_ADDR 和 MASTER_PORT）来初始化分布式环境。在实际应用中，你可能需要根据你的环境来设置这些变量或使用其他初始化方法。

在实际的多GPU环境中，你需要使用 torch.multiprocessing 或其他方式来并行运行多个进程，每个进程调用 main() 函数，并设置不同的 rank。上面的代码示例为了简化，只展示了单个进程的情况。

在多进程环境中，每个进程都会调用 main() 函数，并且每个进程都会看到自己的 rank 和 world_size。每个进程都会在其分配的GPU上运行模型的副本。

由于 DDP 会自动处理数据的分发和梯度的聚合，因此你不需要在代码中显式地处理这些操作。