登录/ 注册
 找回密码
 立即注册
关闭

CSDN热搜

程序园 精品问答 技术交流 资源下载
返回列表 发新帖
首页 › 业界区 › 业界 ›手把手教你多卡分布训练Accelerate使用配置教程 ...

手把手教你多卡分布训练Accelerate使用配置教程

精滂软 昨天 14:31
作者：SkyXZ
CSDN：SkyXZ～-CSDN博客
博客园：SkyXZ - 博客园
开发机环境：Ubuntu 22.04 | 112x CPU | 1TB RAM | 8×NVIDIA A100-40GB | CUDA 12.4 |


  • HuggingFace官方accelerate库中文文档：https://hugging-face.cn/docs/accelerate/index
        在大模型训练和推理任务中，单卡资源往往难以满足性能需求，尤其在模型参数量动辄数亿甚至数十亿的今天，高效利用多卡资源成为提升训练效率的关键手段，目前PyTorch有如下几个常见的多卡并行训练的方式，这几种方式的优缺点都非常明显，那么有没有一个用起简单顺手，计算资源调度又非常优秀的框架呢？
并行方式优点局限性torch.nn.DataParallel使用简单，无需修改太多代码已不推荐使用，存在性能瓶颈，通信开销大，只在主卡做前向/反向调度torch.nn.parallel.DistributedDataParallel（DDP）官方推荐，性能优秀，支持多机多卡配置繁琐，需手动设置启动参数、分布式环境变量、训练脚本需精细改写手动使用 torch.distributed.launch 或 torchrun灵活性强，适合大规模集群易踩坑，调试困难，对新手不友好        有的有的兄弟，那就是HuggingFace推出的Accelerate库，它屏蔽了底层的分布式细节，极大简化了多GPU甚至多机训练的配置和使用流程，适用于 PyTorch 生态下绝大多数任务，本教程将基于本人在A100集群上的实际部署经验，详尽讲解 Accelerate 的配置方式、训练脚本改写建议，以及与 DeepSpeed 等工具的集成策略，帮助你快速构建一个稳定、高效、可复现的多卡分布式训练环境。
Accelerate的安装配置

        Accelerate库具备非常好的跨平台适应能力，支持多种类型的集群环境。在 CPU 训练 方面，它不仅支持单节点多进程的 CPU 分布式训练，也同样支持多节点 CPU 集群训练；在 GPU 和 TPU 环境中，Accelerate 同样支持单节点多卡与多节点多卡的分布式训练，能够灵活适配从本地开发到大规模集群部署的各种场景，其安装也非常的简单，仅需在Python3.8+以上的环境运行一下命令即可完成：
  1. # pip安装
  2. pip install accelerate
  3. # conda安装
  4. conda install -c conda-forge accelerate
复制代码
        在完成安装之后我们需要配置 Accelerate 以适应当前系统的训练设置，完成配置仅需要输入以下命令并按照要求回答问题即可：
  1. # 输入以下命令
  2. accelerate config
  3. # 问题参考如下：
  4. (xq) qi.xiong@instance-qih2207m:~$ accelerate config
  5. --------------------------------------------------------------------------------------------------------------------------------
  6. In which compute environment are you running?
  7. Please select a choice using the arrow or number keys, and selecting with enter
  8. ➔  This machine                                                                                                                      AWS (Amazon SageMaker)   
  9. Which type of machine are you using?                                                                                             Please select a choice using the arrow or number keys, and selecting with enter
  10.     No distributed training
  11.     multi-CPU
  12.     multi-XPU
  13.     multi-HPU
  14. ➔ multi-GPU
  15.     multi-NPU
  16.     multi-MLU
  17.     multi-SDAA
  18.     multi-MUSA
  19.     TPU
  20. How many different machines will you use (use more than 1 for multi-node training)? [1]:
  21. Should distributed operations be checked while running for errors? This can avoid timeout issues but will be slower. [yes/NO]:
  22. Do you wish to optimize your script with torch dynamo?[yes/NO]:
  23. Which dynamo backend would you like to use?                                                                                     Please select a choice using the arrow or number keys, and selecting with enter                                                         eager                                                                                                                                      aot_eager                                                                                                                   ➔  inductor                                                                                                                             aot_ts_nvfuser                                                                                                                        nvprims_nvfuser
  24.     cudagraphs
  25.     ofi
  26.     fx2trt
  27.     onnxrt
  28.     tensorrt
  29.     aot_torchxla_trace_once
  30.     torhchxla_trace_once
  31.     ipex
  32.     tvm
  33. Do you want to customize the defaults sent to torch.compile? [yes/NO]:
  34. Do you want to use DeepSpeed? [yes/NO]:
  35. What GPU(s) (by id) should be used for training on this machine as a comma-separated list? [all]:
  36. Do you wish to use mixed precision?                                                                                             Please select a choice using the arrow or number keys, and selecting with enter                                                  ➔  no                                                                                                                                   fp16                                                                                                                                 bf16                                                                                                                                 fp8                                                                                                                           
复制代码
        完成了如上问题的选择之后便会在以下地址/home/qi.xiong/.cache/huggingface/accelerate/default_config.yaml生成当前的训练配置文件，具体内容格式如下：
  1. compute_environment: LOCAL_MACHINE
  2. debug: true
  3. distributed_type: MULTI_GPU
  4. downcast_bf16: 'no'
  5. dynamo_config:
  6.   dynamo_backend: INDUCTOR
  7. enable_cpu_affinity: false
  8. gpu_ids: 2,3,4,5
  9. machine_rank: 0
  10. main_training_function: main
  11. mixed_precision: 'no'
  12. num_machines: 1
  13. num_processes: 4
  14. rdzv_backend: static
  15. same_network: true
  16. tpu_env: []
  17. tpu_use_cluster: false
  18. tpu_use_sudo: false
  19. use_cpu: false
复制代码
        如果要检查配置的环境是否正确的话我们还可以在命令行运行如下命令：
  1. # 检查配置：
  2. accelerate env
  3. # 示例输出：
  4. (xq) qi.xiong@instance-qih2207m:~$ accelerate env
  5. Copy-and-paste the text below in your GitHub issue
  6. - `Accelerate` version: 1.9.0
  7. - Platform: Linux-5.15.0-72-generic-x86_64-with-glibc2.35
  8. - `accelerate` bash location: /home/qi.xiong/miniconda3/envs/xq/bin/accelerate
  9. - Python version: 3.10.18
  10. - Numpy version: 1.23.4
  11. - PyTorch version: 2.2.2+cu121
  12. - PyTorch accelerator: CUDA
  13. - System RAM: 1005.77 GB
  14. - GPU type: NVIDIA GeForce RTX 4090
  15. - `Accelerate` default config:
  16.         - compute_environment: LOCAL_MACHINE
  17.         - distributed_type: MULTI_GPU
  18.         - mixed_precision: no
  19.         - use_cpu: False
  20.         - debug: True
  21.         - num_processes: 4
  22.         - machine_rank: 0
  23.         - num_machines: 1
  24.         - gpu_ids: 2,3,4,5
  25.         - rdzv_backend: static
  26.         - same_network: True
  27.         - main_training_function: main
  28.         - enable_cpu_affinity: False
  29.         - downcast_bf16: no
  30.         - tpu_use_cluster: False
  31.         - tpu_use_sudo: False
  32.         - tpu_env: []
  33.         - dynamo_config: {'dynamo_backend': 'INDUCTOR'}
复制代码
        同时我们还可以运行以下命令来对分布式配置进行一次测试来验证我们的Accelerate是否已在服务器上正确配置并运行，其中的arguments为我们上述配置文件生成的路径，默认情况下会指向我们前面在.cache下生成的default_config.yaml
  1. # 测试命令
  2. accelerate test [arguments]
  3. options:
  4.   -h, --help            show this help message and exit
  5.   --config_file CONFIG_FILE, --config-file CONFIG_FILE
  6. The path to use to store the config file. Will default to a file named default_config.yaml in the cache location, which is the
  7. content of the environment `HF_HOME` suffixed with 'accelerate', or if you don't have such an environment variable, your cache
  8. directory ('~/.cache' or the content of `XDG_CACHE_HOME`) suffixed with 'huggingface'.
  9. # 出现以下log即代表配置正确无误
  10. stdout: **Breakpoint trigger test**s
  11. Test is a success! You are ready for your distributed training!
复制代码
完成训练代码的配置

        在完成了基本配置之后我们便可以开始丢弃过时的nn.DataParallel修改我们的训练脚本以支持Accelerate啦，我们以我之前自己搭的一个ViT训练脚本为例子，这个脚本使用的是典型的nn.DataParallel分布式训练方式，通过将模型包裹在 nn.DataParallel(model) 中，并手动选择可用 GPU，这个方式有几个缺点首先是nn.DataParallel 只能在单节点使用，且依赖主 GPU 聚合梯度，存在性能瓶颈；其次它不支持多进程训练，对模型部署和迁移都不够灵活，所以接下来我们来对这个训练脚本进行升级，原始训练脚本如下：
  1. from model.transformer_net import VisonTransformer
  2. from cifar10_dataset import CIFAR10DataLoad
  3. import os
  4. import torch
  5. import numpy as np
  6. import torch.nn as nn
  7. import torch.optim as optim
  8. from tqdm import tqdm
  9. import matplotlib.pyplot as plt
  11. os.environ["CUDA_VISIBLE_DEVICES"] = "4,5"
  13. def train_cifar10(
  14.     root="/home/qi.xiong/Dataset/cifar-10-batches-py", # CIFAR-10数据集根目录
  15.     img_size=224,
  16.     patch_size=16,      # 改为4，这样有64个patch (32/4)²
  17.     in_channels=3,
  18.     num_features=768,  # 减少特征维度，适应较小的模型
  19.     depth=12,          # 减少层数
  20.     num_heads=12,      # 减少注意力头数
  21.     mlp_ratio=4.0,
  22.     qkv_bias=True,
  23.     drop_rate=0.1,
  24.     attn_drop_rate=0.1,
  25.     drop_path_rate=0.1,
  26.     epochs=20,
  27.     batch_size=128,
  28.     num_workers=8,
  29.     lr=1e-4,
  30.     device="cuda"
  31. ):
  32.     device = device or ("cuda" if torch.cuda.is_available() else "cpu")
  33.     print(f"使用设备: {device}")
  34.     # 数据加载
  35.     train_loader, test_loader = CIFAR10DataLoad(root, batch_size, num_workers, img_size)
  36.     num_classes = 10  # CIFAR-10固定10个类别
  37.     input_shape = (in_channels, img_size, img_size)
  39.     model = VisonTransformer(
  40.         input_shape=input_shape,
  41.         patch_size=patch_size,
  42.         in_channels=in_channels,
  43.         num_classes=num_classes,
  44.         num_features=num_features,
  45.         depth=depth,
  46.         num_heads=num_heads,
  47.         mlp_ratio=mlp_ratio,
  48.         qkv_bias=qkv_bias,
  49.         drop_rate=drop_rate,
  50.         attn_drop_rate=attn_drop_rate,
  51.         drop_path_rate=drop_path_rate,
  52.         act_layer=nn.GELU,
  53.         norm_layer=nn.LayerNorm
  54.     ).to(device)
  55.    
  56.     if torch.cuda.device_count() > 1:
  57.         print(f"使用 {torch.cuda.device_count()} 张GPU进行训练")
  58.         model = nn.DataParallel(model)
  60.     print(f"模型参数量: {sum(p.numel() for p in model.parameters() if p.requires_grad):,}")
  62.     criterion = nn.CrossEntropyLoss()
  63.     optimizer = optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)
  64.     scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
  65.    
  66.     best_acc = 0
  67.     train_loss_list, test_loss_list = [], []
  68.     train_acc_list, test_acc_list = [], []
  70.     for epoch in range(epochs):
  71.         model.train()
  72.         total_loss, correct, total = 0, 0, 0
  73.         pbar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}")
  74.         
  75.         for images, labels in pbar:
  76.             images, labels = images.to(device), labels.to(device)
  77.             optimizer.zero_grad()
  78.             outputs = model(images)
  79.             loss = criterion(outputs, labels)
  80.             loss.backward()
  81.             optimizer.step()
  82.             
  83.             total_loss += loss.item() * images.size(0)
  84.             _, preds = outputs.max(1)
  85.             correct += preds.eq(labels).sum().item()
  86.             total += labels.size(0)
  87.             
  88.             # 更新进度条
  89.             pbar.set_postfix({
  90.                 'Loss': f'{loss.item():.4f}',
  91.                 'Acc': f'{100.*correct/total:.2f}%'
  92.             })
  93.         
  94.         train_loss = total_loss / total
  95.         train_acc = correct / total
  96.         train_loss_list.append(train_loss)
  97.         train_acc_list.append(train_acc)
  99.         # 测试阶段
  100.         model.eval()
  101.         test_loss, test_correct, test_total = 0, 0, 0
  102.         with torch.no_grad():
  103.             for images, labels in test_loader:
  104.                 images, labels = images.to(device), labels.to(device)
  105.                 outputs = model(images)
  106.                 loss = criterion(outputs, labels)
  107.                 test_loss += loss.item() * images.size(0)
  108.                 _, preds = outputs.max(1)
  109.                 test_correct += preds.eq(labels).sum().item()
  110.                 test_total += labels.size(0)
  111.         
  112.         test_loss = test_loss / test_total
  113.         test_acc = test_correct / test_total
  114.         test_loss_list.append(test_loss)
  115.         test_acc_list.append(test_acc)
  116.         
  117.         # 更新学习率
  118.         scheduler.step()
  119.         current_lr = scheduler.get_last_lr()[0]
  121.         print(f"Epoch {epoch+1}: Train Loss={train_loss:.4f}, Train Acc={train_acc:.4f}, "
  122.               f"Test Loss={test_loss:.4f}, Test Acc={test_acc:.4f}, LR={current_lr:.6f}")
  123.         
  124.         # 保存最优模型
  125.         if test_acc > best_acc:
  126.             best_acc = test_acc
  127.             torch.save(model.module.state_dict(), "best_vit_cifar10.pth")
  128.             print(f"保存最优模型，测试准确率: {best_acc:.4f}")
  130.     # 可视化loss和acc
  131.     plt.figure(figsize=(12, 4))
  132.    
  133.     plt.subplot(1, 2, 1)
  134.     plt.plot(train_loss_list, label="Train Loss")
  135.     plt.plot(test_loss_list, label="Test Loss")
  136.     plt.legend()
  137.     plt.title("Loss Curve")
  138.     plt.xlabel("Epoch")
  139.     plt.ylabel("Loss")
  140.    
  141.     plt.subplot(1, 2, 2)
  142.     plt.plot(train_acc_list, label="Train Acc")
  143.     plt.plot(test_acc_list, label="Test Acc")
  144.     plt.legend()
  145.     plt.title("Accuracy Curve")
  146.     plt.xlabel("Epoch")
  147.     plt.ylabel("Accuracy")
  148.    
  149.     plt.tight_layout()
  150.     plt.savefig("cifar10_training_curves.png", dpi=300, bbox_inches='tight')
  151.     plt.show()
  152.    
  153.     print(f"训练完成！最优测试准确率: {best_acc:.4f}")
  154.     return model, best_acc
  156. if __name__ == "__main__":
  157.     train_cifar10()
复制代码
        我们首先要做的便是导入Accelerate库，同时继承Accelerator()类，并把我们最开始的显式指定GPU以及device设置给删去让accelerate自动帮我们进行配置
  1. from model.transformer_net import VisonTransformer
  2. from cifar10_dataset import CIFAR10DataLoad
  3. import os
  4. import torch
  5. import numpy as np
  6. import torch.nn as nn
  7. import torch.optim as optim
  8. from tqdm import tqdm
  9. import matplotlib.pyplot as plt
  10. # 导入accelerate
  11. from accelerate import Accelerator
  12. # 删去以下手动GPU设置，让Accelerate进行管理
  13. # os.environ["CUDA_VISIBLE_DEVICES"] = "4,5"
  15. def train_cifar10(
  16.     root="/home/qi.xiong/Dataset/cifar-10-batches-py", # CIFAR-10数据集根目录
  17.     img_size=224,
  18.     patch_size=16,      # 改为4，这样有64个patch (32/4)²
  19.     in_channels=3,
  20.     num_features=768,  # 减少特征维度，适应较小的模型
  21.     depth=12,          # 减少层数
  22.     num_heads=12,      # 减少注意力头数
  23.     mlp_ratio=4.0,
  24.     qkv_bias=True,
  25.     drop_rate=0.1,
  26.     attn_drop_rate=0.1,
  27.     drop_path_rate=0.1,
  28.     epochs=20,
  29.     batch_size=128,
  30.     num_workers=8,
  31.     lr=1e-4,
  32.     device="cuda"
  33. ):
  34.     # 注释掉以下device设置内容，并初始化Accelerator()类，并使用accelerator.device来管理设备
  35.     # device = device or ("cuda" if torch.cuda.is_available() else "cpu")
  36.     accelerator = Accelerator()
  37.     device = accelerator.device
复制代码
        我们在上面已经使用了accelerator.device，因此我们需要删去代码中所有的.to(device)操作，并不用nn.DataParallel包裹我们的模型
1.png

        接下来我们使用Accelerate来包装我们的模型、优化器和数据加载器，我们用到的是Accelerate类中的prepare方法（具体介绍见：https://hugging-face.cn/docs/accelerate/v1.6.0/en/package_reference/accelerator#accelerate.Accelerator.prepare），这个方法也是HF官方最推荐的方法，可以一步准备多个组件，其官方介绍如下，我们可以传入任何顺序组合的torch.utils.data.DataLoader数据加载器、torch.nn.Module模型、torch.optim.Optimizer优化器以及torch.optim.lr_scheduler.LRScheduler学习器，同时我们还可以使用device_placement来选择是否让 Accelerate 自动将这些对象转到目标设备，如果仅想让某几个Args转移，可以传一个与 args 相同长度的布尔列表例如 [True, True, False, True]）：
  1. def prepare(self, *args, device_placement=None):
  2.     """
  3.     Prepare all objects passed in `args` for distributed training and mixed precision, then return them in the same
  4.     order.
  5.     Args:
  6.         *args (list of objects):
  7.             Any of the following type of objects:
  9.             - `torch.utils.data.DataLoader`: PyTorch Dataloader
  10.             - `torch.nn.Module`: PyTorch Module
  11.             - `torch.optim.Optimizer`: PyTorch Optimizer
  12.             - `torch.optim.lr_scheduler.LRScheduler`: PyTorch LR Scheduler
  14.         device_placement (`list[bool]`, *optional*):
  15.             Used to customize whether automatic device placement should be performed for each object passed. Needs
  16.             to be a list of the same length as `args`. Not compatible with DeepSpeed or FSDP.
  17.     <Tip>
  18.       You don't need to prepare a model if you only use it for inference without any kind of mixed precision
  19.     </Tip>
  20.     """
  21. # device_placement example:
  22. device_placement = [True, True, False, False]
  23. model, optimizer, data_loader, scheduler = accelerator.prepare(
  24.     model, optimizer, data_loader, scheduler, device_placement=device_placement
  25. )
复制代码
        虽然 accelerator.prepare() 是最推荐用法，但 Accelerate 也提供了以下更细化的封装函数，适用于只想准备某一类对象的场景，它们的功能其实和 prepare() 一致，只不过拆解得更细，灵活性更高，下面是简要介绍（实际中使用场景较少，因此不做深入讲解）：
  1. # 1.prepare_data_loader方法只包装一个数据加载器DataLoader
  2. # 用于将单个数据加载器适配到分布式环境下，并可选择是否自动将数据移动到目标设备（如 GPU）
  3. import torch
  4. from accelerate import Accelerato
  5. accelerator = Accelerator()
  6. data_loader = torch.utils.data.DataLoader(...)
  7. data_loader = accelerator.prepare_data_loader(data_loader, device_placement=True)
  8.         # 参数说明：
  9.             device_placement=True：是否自动将 batch 移动到正确设备（默认使用 accelerator.device）
  10.         slice_fn_for_dispatch：用于多进程数据切分，一般无需手动设置
  11. # 2.prepare_model方法只准备模型
  12. # 用于将模型加载到合适设备，并根据设置（如是否开启混合精度、是否启用 torch.compile）对其进行优化。
  13. from accelerate import Accelerator
  14. accelerator = Accelerator()
  15. model = accelerator.prepare_model(model, device_placement=True, evaluation_mode=False)
  16.         # 参数说明：
  17.             device_placement：是否将模型移动到正确设备
  18.         evaluation_mode：如果设置为 True，表示模型仅用于推理，不参与训练，适用于纯推理场景
  19. # 3.prepare_optimizer方法只包装优化器
  20. # 适配优化器在多卡训练场景下的兼容性，确保其工作在正确的设备上
  21. import torch
  22. from accelerate import Accelerator
  23. accelerator = Accelerator()
  24. optimizer = torch.optim.Adam(...)
  25. optimizer = accelerator.prepare_optimizer(optimizer, device_placement=True)
  26.         # 参数说明：
  27.             device_placement：是否将模型移动到正确设备
  28. # 4.prepare_scheduler只包装学习率调度器
  29. # 适配调度器逻辑到分布式环境，通常在 optimizer 已被 prepare 后调用
  30. import torch
  31. from accelerate import Accelerator
  32. accelerator = Accelerator()
  33. optimizer = torch.optim.Adam(...)
  34. scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, ...)
  35. scheduler = accelerator.prepare_scheduler(scheduler)
复制代码
        因此我们这里使用推荐的prepare方法，修改后的代码如下：
  1. # 具体修改如下
  2. criterion = nn.CrossEntropyLoss()
  3. optimizer = optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)
  4. scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
  5. # 使用Accelerate包装我们的模型、优化器和数据加载器，并默认全部加载到目标设备也就是GPU上
  6. model, optimizer, train_loader, test_loader, scheduler = accelerator.prepare(
  7.     model, optimizer, train_loader, test_loader, scheduler
  8. )
复制代码
        接着我们删去所有的images, labels = images.to(device), labels.to(device)，用accelerator.backward(loss)来替换原先的loss.backward()，同时将模型的保存修改为以下：
  1. # 模型保存，修改以下为：
  2. # torch.save(model.module.state_dict(), "best_vit_cifar10.pth")
  3. accelerator.save_state("best_vit_cifar10")
复制代码
        至此我们accelerate的适配就基本完成啦！接着我们在添加一点小细节，比如训练的进度显示及进程同步，在分布式训练环境中，多个进程同时运行会导致进度条显示混乱和日志重复打印因此我们需要使用accelerate给我们提供的更简洁、安全，会自动只在主进程输出的print方法，仅需将原始的print修改为accelerate.print即可，同时要注意的是我们这里的batch_size将会变为每张卡上加载的大小，如果需要自动划分的话自行除以线程数即可，修改后的完整代码如下：
  1. from model.transformer_net import VisonTransformer
  2. from cifar10_dataset import CIFAR10DataLoad
  3. import os
  4. import torch
  5. import numpy as np
  6. import torch.nn as nn
  7. import torch.optim as optim
  8. from tqdm import tqdm
  9. import matplotlib.pyplot as plt
  10. from accelerate import Accelerator
  11. from accelerate.logging import get_logger
  12. # 移除手动GPU设置，让Accelerate管理
  13. # os.environ["CUDA_VISIBLE_DEVICES"] = "4,5"
  14. def train_cifar10(
  15.     root="/home/qi.xiong/Dataset/cifar-10-batches-py", # CIFAR-10数据集根目录
  16.     img_size=224,
  17.     patch_size=16,      
  18.     in_channels=3,
  19.     num_features=768,
  20.     depth=12,         
  21.     num_heads=12,   
  22.     mlp_ratio=4.0,
  23.     qkv_bias=True,
  24.     drop_rate=0.1,
  25.     attn_drop_rate=0.1,
  26.     drop_path_rate=0.1,
  27.     epochs=20,
  28.     batch_size=128,
  29.     num_workers=8,
  30.     lr=1e-4,
  31.     device="cuda"
  32. ):
  33.     # 初始化Accelerate
  34.     accelerator = Accelerator()
  35.     device = accelerator.device
  36.     # 使用 accelerator.print() 自动只在主进程输出
  37.     accelerator.print(f"使用设备: {device}")
  38.     accelerator.print(f"分布式训练: {accelerator.num_processes} 个进程")
  39.     # 数据加载
  40.     train_loader, test_loader = CIFAR10DataLoad(root, batch_size, num_workers, img_size)
  41.     num_classes = 10  # CIFAR-10固定10个类别
  42.     input_shape = (in_channels, img_size, img_size)
  43.     accelerator.print(f"训练数据: {len(train_loader.dataset)} 张图像")
  44.     accelerator.print(f"测试数据: {len(test_loader.dataset)} 张图像")
  45.     accelerator.print(f"批次大小: {batch_size}, 训练步数: {len(train_loader)}")
  46.     # 模型
  47.     model = VisonTransformer(
  48.         input_shape=input_shape,
  49.         patch_size=patch_size,
  50.         in_channels=in_channels,
  51.         num_classes=num_classes,
  52.         num_features=num_features,
  53.         depth=depth,
  54.         num_heads=num_heads,
  55.         mlp_ratio=mlp_ratio,
  56.         qkv_bias=qkv_bias,
  57.         drop_rate=drop_rate,
  58.         attn_drop_rate=attn_drop_rate,
  59.         drop_path_rate=drop_path_rate,
  60.         act_layer=nn.GELU,
  61.         norm_layer=nn.LayerNorm
  62.     )
  63.     accelerator.print(f"模型参数量: {sum(p.numel() for p in model.parameters() if p.requires_grad):,}")
  64.     criterion = nn.CrossEntropyLoss()
  65.     optimizer = optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)
  66.     scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
  67.     # 使用Accelerate包装模型、优化器和数据加载器
  68.     model, optimizer, train_loader, test_loader, scheduler = accelerator.prepare(
  69.         model, optimizer, train_loader, test_loader, scheduler
  70.     )
  71.     best_acc = 0
  72.     train_loss_list, test_loss_list = [], []
  73.     train_acc_list, test_acc_list = [], []
  74.     for epoch in range(epochs):
  75.         # 训练阶段
  76.         model.train()
  77.         total_loss, correct, total = 0, 0, 0
  78.         
  79.         if accelerator.is_local_main_process:
  80.             pbar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}")
  81.         else:
  82.             pbar = train_loader
  83.         for images, labels in pbar:
  84.             optimizer.zero_grad()
  85.             outputs = model(images)
  86.             loss = criterion(outputs, labels)
  87.             accelerator.backward(loss)
  88.             optimizer.step()
  89.             _, preds = outputs.max(1)
  90.             gathered_loss = accelerator.gather(loss.unsqueeze(0))
  91.             gathered_correct = accelerator.gather(preds.eq(labels).sum().unsqueeze(0))
  92.             gathered_total = accelerator.gather(torch.tensor(labels.size(0)).unsqueeze(0).to(device))
  93.             global_loss = gathered_loss.mean().item()
  94.             global_correct = gathered_correct.sum().item()
  95.             global_total = gathered_total.sum().item()
  96.             global_acc = 100. * global_correct / global_total if global_total > 0 else 0
  97.             total_loss += loss.item() * images.size(0)
  98.             correct += preds.eq(labels).sum().item()
  99.             total += labels.size(0)
  100.             if accelerator.is_local_main_process and hasattr(pbar, 'set_postfix'):
  101.                 pbar.set_postfix({
  102.                     'Loss': f'{global_loss:.4f}',
  103.                     'Acc': f'{global_acc:.2f}%'
  104.                 })
  105.         train_loss = total_loss / total
  106.         train_acc = correct / total
  107.         train_loss_list.append(train_loss)
  108.         train_acc_list.append(train_acc)
  109.         # 测试阶段
  110.         model.eval()
  111.         test_loss, test_correct, test_total = 0, 0, 0
  112.         with torch.no_grad():
  113.             for images, labels in test_loader:
  114.                 outputs = model(images)
  115.                 loss = criterion(outputs, labels)
  116.                 test_loss += loss.item() * images.size(0)
  117.                 _, preds = outputs.max(1)
  118.                 test_correct += preds.eq(labels).sum().item()
  119.                 test_total += labels.size(0)
  120.         
  121.         test_loss = test_loss / test_total
  122.         test_acc = test_correct / test_total
  123.         test_loss_list.append(test_loss)
  124.         test_acc_list.append(test_acc)
  125.         # 更新学习率
  126.         scheduler.step()
  127.         current_lr = scheduler.get_last_lr()[0]
  128.         accelerator.print(f"Epoch {epoch+1}: Train Loss={train_loss:.4f}, Train Acc={train_acc:.4f}, "
  129.                   f"Test Loss={test_loss:.4f}, Test Acc={test_acc:.4f}, LR={current_lr:.6f}")
  130.         
  131.         accelerator.wait_for_everyone()
  132.         if test_acc > best_acc:
  133.             best_acc = test_acc
  134.             accelerator.save_state("best_vit_cifar10")
  135.             accelerator.print(f"保存最优模型，测试准确率: {best_acc:.4f}")
  136.     # 可视化loss和acc
  137.     plt.figure(figsize=(12, 4)
  138.     plt.subplot(1, 2, 1)
  139.     plt.plot(train_loss_list, label="Train Loss")
  140.     plt.plot(test_loss_list, label="Test Loss")
  141.     plt.legend()
  142.     plt.title("Loss Curve")
  143.     plt.xlabel("Epoch")
  144.     plt.ylabel("Loss")
  145.    
  146.     plt.subplot(1, 2, 2)
  147.     plt.plot(train_acc_list, label="Train Acc")
  148.     plt.plot(test_acc_list, label="Test Acc")
  149.     plt.legend()
  150.     plt.title("Accuracy Curve")
  151.     plt.xlabel("Epoch")
  152.     plt.ylabel("Accuracy")
  153.     accelerator.wait_for_everyone()
  154.     if accelerator.is_local_main_process:
  155.         plt.tight_layout()
  156.         plt.savefig("cifar10_training_curves.png", dpi=300, bbox_inches='tight')
  157.         plt.show()
  158.         accelerator.print(f"训练完成！最优测试准确率: {best_acc:.4f}")
  159.     return model, best_acc
  160. if __name__ == "__main__":
  161.     train_cifar10()
复制代码
        完成了训练之后我们便可以使用以下命令开启分布式训练啦！由于我们前面已经使用accelerate config对训练进行了配置因此我们直接启动即可，如果需要使用自定义或者其他位置的配置文件的话手动传入配置文件位置即可
  1. accelerate launch --multi_gpu train.py
复制代码
2.png

        当然，accelerate也支持不提前配置而是直接启动传入参数，具体有哪些参数可以使用大家可以看官方文档启动 Accelerate 脚本 - Hugging Face 机器学习平台或者是使用如下命令
  1. accelerate launch -h
复制代码
来源：程序园用户自行投稿发布，如果侵权，请联系站长删除
免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！

使用道具 举报
B Color Image Link Quote Code Smilies
您需要登录后才可以回帖 登录 | 立即注册
发帖
精滂软
昨天 14:31

0

粉丝关注

14

主题发布

板块介绍填写区域，请于后台编辑
微信扫一扫