• 多卡服务器使用

    参考文章：https://blog.csdn.net/qq_42255269/article/details/123427094?spm=1001.2014.3001.5506

    示例程序

    即最简单的用 vgg16模型 对 CIFAR10数据集 进行10分类

    from torch import nn
import torchvision
import argparse
from torch.utils.data import DataLoader
import torch


class MyNet(nn.Module):
    def __init__(self):
        super(MyNet, self).__init__()
        self.vgg16 = torchvision.models.vgg16()
        self.fc = nn.Linear(1000, 10)

    def forward(self, x):
        out = self.vgg16(x)
        out = self.fc(out)
        return out


if __name__ == '__main__':

    # 获取batch_size参数
    parser = argparse.ArgumentParser()
    parser.add_argument('--batch_size', type=int)
    args = parser.parse_args()
    
    device = 'cuda'

    # 配置数据集
    train_data = torchvision.datasets.CIFAR10(root="data", train=True, transform=torchvision.transforms.ToTensor(), download=True)
    train_dataloader = DataLoader(train_data, batch_size=args.batch_size)

    # 创建网络模型
    net = MyNet().to(device=device)
    # 损失函数
    loss_fn = nn.CrossEntropyLoss()
    # 优化器
    optimizer = torch.optim.SGD(net.parameters(), lr=1e-2)

    for i in range(1000):
        for step, data in enumerate(train_dataloader):
            imgs, targets = data
            imgs = imgs.to(device=device)
            targets = targets.to(device=device)
            outputs = net(imgs)
            loss = loss_fn(outputs, targets)

            # 优化器优化模型
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            print("epoch:{} step:{} loss:{}".format(i, step, loss.item()))


    • 1
    • 2
    • 3
    • 4
    • 5
    • 6
    • 7
    • 8
    • 9
    • 10
    • 11
    • 12
    • 13
    • 14
    • 15
    • 16
    • 17
    • 18
    • 19
    • 20
    • 21
    • 22
    • 23
    • 24
    • 25
    • 26
    • 27
    • 28
    • 29
    • 30
    • 31
    • 32
    • 33
    • 34
    • 35
    • 36
    • 37
    • 38
    • 39
    • 40
    • 41
    • 42
    • 43
    • 44
    • 45
    • 46
    • 47
    • 48
    • 49
    • 50
    • 51
    • 52
    • 53
    • 54

    单 卡机器 单 卡运行

    python main.py --batch_size 16

    • 1

    多 卡机器 单 卡运行

    # 卡0运行
CUDA_VISIBLE_DEVICES=0 python main.py --batch_size 16
# 卡1运行
CUDA_VISIBLE_DEVICES=1 python main.py --batch_size 16

    • 1
    • 2
    • 3
    • 4

    多 卡机器 多 卡运行

    nn.DataParallel（不推荐）

    优点：使用简单，对代码的修改比较小（仅3处修改）

    缺点：只是提高训练速度但不能提高batch_size（数据需要先加载到主GPU上再进行其他操作，所有主GPU的显存限制了batch_size的大小）

    # batch_size大小应为显卡数量的整数倍
python main.py --batch_size 512

    • 1
    • 2
    from torch import nn
import torchvision
import argparse
from torch.utils.data import DataLoader
import torch


class MyNet(nn.Module):
    def __init__(self):
        super(MyNet, self).__init__()
        self.vgg16 = torchvision.models.vgg16()
        self.fc = nn.Linear(1000, 10)

    def forward(self, x):
        out = self.vgg16(x)
        out = self.fc(out)
        return out


if __name__ == '__main__':

    # 获取batch_size参数
    parser = argparse.ArgumentParser()
    parser.add_argument('--batch_size', type=int)
    args = parser.parse_args()
    
    # 修改1
    gpus = [0, 1]
    device = 'cuda:{}'.format(gpus[0])
    
    # 配置数据集
    train_data = torchvision.datasets.CIFAR10(root="data", train=True, transform=torchvision.transforms.ToTensor(), download=True)
    train_dataloader = DataLoader(train_data, batch_size=args.batch_size)

    # 修改2 创建网络模型
    net = nn.DataParallel(MyNet().to(device=device), device_ids=gpus, output_device=gpus[0])
    # 损失函数
    loss_fn = nn.CrossEntropyLoss()
    # 优化器
    optimizer = torch.optim.SGD(net.parameters(), lr=1e-2)

    for i in range(1000):
        for step, data in enumerate(train_dataloader):
            imgs, targets = data
            # 修改3
            imgs = imgs.to(device=device, non_blocking=True)
            targets = targets.to(device=device, non_blocking=True)
            outputs = net(imgs)
            loss = loss_fn(outputs, targets)

            # 优化器优化模型
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            print("epoch:{} step:{} loss:{}".format(i, step, loss.item()))


    • 1
    • 2
    • 3
    • 4
    • 5
    • 6
    • 7
    • 8
    • 9
    • 10
    • 11
    • 12
    • 13
    • 14
    • 15
    • 16
    • 17
    • 18
    • 19
    • 20
    • 21
    • 22
    • 23
    • 24
    • 25
    • 26
    • 27
    • 28
    • 29
    • 30
    • 31
    • 32
    • 33
    • 34
    • 35
    • 36
    • 37
    • 38
    • 39
    • 40
    • 41
    • 42
    • 43
    • 44
    • 45
    • 46
    • 47
    • 48
    • 49
    • 50
    • 51
    • 52
    • 53
    • 54
    • 55
    • 56
    • 57

    torch.distributed（推荐）

    优点：可以提高训练速度和batch_size

    缺点：使用稍复杂，数据集不能shuffle

    新版本

    CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.run --nproc_per_node=2 main.py --batch_size 512

    • 1
    from torch import nn
import torchvision
import argparse
from torch.utils.data import DataLoader
import torch
import torch.distributed as dist
import os


class MyNet(nn.Module):
    def __init__(self):
        super(MyNet, self).__init__()
        self.vgg16 = torchvision.models.vgg16()
        self.fc = nn.Linear(1000, 10)

    def forward(self, x):
        out = self.vgg16(x)
        out = self.fc(out)
        return out


if __name__ == '__main__':

    # 获取batch_size参数
    parser = argparse.ArgumentParser()
    parser.add_argument('--batch_size', type=int)
    args = parser.parse_args()

    # 修改1
    local_rank = int(os.environ["LOCAL_RANK"])

    # 修改2 设置GPU之间通信使用的后端和端口
    dist.init_process_group(backend='nccl')

    # 修改3
    device = 'cuda:{}'.format(local_rank)

    # 配置数据集
    train_data = torchvision.datasets.CIFAR10(root="data", train=True, transform=torchvision.transforms.ToTensor(), download=True)
    # 修改4 注意不要设置DataLoader的shuffle为True
    train_sampler = torch.utils.data.distributed.DistributedSampler(train_data)
    train_dataloader = DataLoader(train_data, batch_size=args.batch_size, sampler=train_sampler)

    # 修改5 创建网络模型
    net = nn.parallel.DistributedDataParallel(MyNet().to(device), device_ids=[local_rank])
    # 损失函数
    loss_fn = nn.CrossEntropyLoss()
    # 优化器
    optimizer = torch.optim.SGD(net.parameters(), lr=1e-2)

    for i in range(1000):
        for step, data in enumerate(train_dataloader):
            imgs, targets = data
            # 修改6
            imgs = imgs.to(device=device, non_blocking=True)
            targets = targets.to(device=device, non_blocking=True)
            outputs = net(imgs)
            loss = loss_fn(outputs, targets)

            # 优化器优化模型
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            print("epoch:{} step:{} loss:{}".format(i, step, loss.item()))


    • 1
    • 2
    • 3
    • 4
    • 5
    • 6
    • 7
    • 8
    • 9
    • 10
    • 11
    • 12
    • 13
    • 14
    • 15
    • 16
    • 17
    • 18
    • 19
    • 20
    • 21
    • 22
    • 23
    • 24
    • 25
    • 26
    • 27
    • 28
    • 29
    • 30
    • 31
    • 32
    • 33
    • 34
    • 35
    • 36
    • 37
    • 38
    • 39
    • 40
    • 41
    • 42
    • 43
    • 44
    • 45
    • 46
    • 47
    • 48
    • 49
    • 50
    • 51
    • 52
    • 53
    • 54
    • 55
    • 56
    • 57
    • 58
    • 59
    • 60
    • 61
    • 62
    • 63
    • 64
    • 65
    • 66

    老版本torch.distributed.launch

    CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --nproc_per_node=2 main.py --batch_size 512

    • 1
    from torch import nn
import torchvision
import argparse
from torch.utils.data import DataLoader
import torch
import torch.distributed as dist


class MyNet(nn.Module):
    def __init__(self):
        super(MyNet, self).__init__()
        self.vgg16 = torchvision.models.vgg16()
        self.fc = nn.Linear(1000, 10)

    def forward(self, x):
        out = self.vgg16(x)
        out = self.fc(out)
        return out


if __name__ == '__main__':

    # 获取batch_size参数
    parser = argparse.ArgumentParser()
    parser.add_argument('--batch_size', type=int)
    # 修改1
    parser.add_argument('--local_rank', default=-1, type=int,
                        help='node rank for distributed training')
    args = parser.parse_args()

    # 修改2 设置GPU之间通信使用的后端和端口
    dist.init_process_group(backend='nccl')

    # 修改3
    device = 'cuda:{}'.format(args.local_rank)

    # 配置数据集
    train_data = torchvision.datasets.CIFAR10(root="data", train=True, transform=torchvision.transforms.ToTensor(), download=True)
    # 修改4 注意不要设置DataLoader的shuffle为True
    train_sampler = torch.utils.data.distributed.DistributedSampler(train_data)
    train_dataloader = DataLoader(train_data, batch_size=args.batch_size, sampler=train_sampler)

    # 修改5 创建网络模型
    net = nn.parallel.DistributedDataParallel(MyNet().to(device), device_ids=[args.local_rank])
    # 损失函数
    loss_fn = nn.CrossEntropyLoss()
    # 优化器
    optimizer = torch.optim.SGD(net.parameters(), lr=1e-2)

    for i in range(1000):
        for step, data in enumerate(train_dataloader):
            imgs, targets = data
            # 修改6
            imgs = imgs.to(device=device, non_blocking=True)
            targets = targets.to(device=device, non_blocking=True)
            outputs = net(imgs)
            loss = loss_fn(outputs, targets)

            # 优化器优化模型
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            print("epoch:{} step:{} loss:{}".format(i, step, loss.item()))


    • 1
    • 2
    • 3
    • 4
    • 5
    • 6
    • 7
    • 8
    • 9
    • 10
    • 11
    • 12
    • 13
    • 14
    • 15
    • 16
    • 17
    • 18
    • 19
    • 20
    • 21
    • 22
    • 23
    • 24
    • 25
    • 26
    • 27
    • 28
    • 29
    • 30
    • 31
    • 32
    • 33
    • 34
    • 35
    • 36
    • 37
    • 38
    • 39
    • 40
    • 41
    • 42
    • 43
    • 44
    • 45
    • 46
    • 47
    • 48
    • 49
    • 50
    • 51
    • 52
    • 53
    • 54
    • 55
    • 56
    • 57
    • 58
    • 59
    • 60
    • 61
    • 62
    • 63
    • 64
    • 65
  • 相关阅读:
    不平衡数据处理--学习笔记
    springboot使用redis
    《全职猎人》
    vue大屏可视化项目页面全屏展示（强制缩放实现）
    新时代布局新特性 -- 容器查询
    C++11常用新特性——右值引用&&
    基于Echarts实现可视化数据大屏董事会指标体系层级结构系统
    渲染农场：设计师提高工作效率的得力助手
    不敢相信，Nginx 还能这么玩？
    花998购买的拍摄技巧和7天起号培训文档，学了一周的总结。
  • 原文地址：https://blog.csdn.net/qq_42283621/article/details/125526711