当前位置:   article > 正文

【模型训练】深度学习多GPU训练的不同方式及代码实现_深度学习使用多张卡训练模型

深度学习使用多张卡训练模型


前言

pytorch中,常见的GPU训练方式有:单机单卡训练,单机多卡训练及多机多卡训练。


一、单机单卡训练

单机单卡训练,是最常用的训练方式。相比CPU训练,只需将模型及数据拷贝到cuda中即可。
步骤:
1、使用torch.cuda.is_available()判断本机是否有cuda;
2、将模型复制到cuda中:net.to(device);
3、将训练数据及验证数据复制到cuda中:images.to(device)
完整代码看我的GitHub。

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

net = AlexNet(num_classes=5, init_weights=True)
net.to(device)
loss_function = nn.CrossEntropyLoss()

optimizer = optim.SGD(net.parameters(), lr=0.0002)

epochs = 10
save_path = './AlexNet.pth'
best_acc = 0.0
train_steps = len(train_loader)
for epoch in range(epochs):
   # train
   net.train()
   running_loss = 0.0
   train_bar = tqdm(train_loader, file=sys.stdout)
   for step, data in enumerate(train_bar):
       images, labels = data
       optimizer.zero_grad()
       outputs = net(images.to(device))
       loss = loss_function(outputs, labels.to(device))
       loss.backward()
       optimizer.step()
       running_loss += loss.item()
       train_bar.desc = "train epoch[{}/{}] loss:{:.3f}".format(epoch + 1, epochs, loss)
   
   # validate
   net.eval()
   acc = 0.0  # accumulate accurate number / epoch
   with torch.no_grad():
       val_bar = tqdm(validate_loader, file=sys.stdout)
       for val_data in val_bar:
           val_images, val_labels = val_data
           outputs = net(val_images.to(device))
           predict_y = torch.max(outputs, dim=1)[1]
           acc += torch.eq(predict_y, val_labels.to(device)).sum().item()
   val_accurate = acc / val_num
   print('[epoch %d] train_loss: %.3f  val_accuracy: %.3f' %
         (epoch + 1, running_loss / train_steps, val_accurate))
   if val_accurate > best_acc:
       best_acc = val_accurate
       torch.save(net.state_dict(), save_path)
       
  • 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

二、单机多卡训练

单机多卡训练有两种方式:
(1)torch.nn.DataParallel:DP模式,简单来说,就是一个进程控制多个GPU训练。
(2)torch.nn.parallel.DistributedDataParallel:DDP模式,多个进程控制多个GPU训练。

1.torch.nn.DataParallel

由于该方法采用单进程训练,效率慢,故现在很少采用。
步骤:
1、使用torch.cuda.device_count()判断本机是否有多个GPU,是否支持多GPU进行训练;
2、将模型使用torch.nn.DataParallel()进行包装,要指定device_ids用来训练的GPU卡;
3、将训练数据传入cuda()中
4、模型保存时,注意要使用torch.save(model.module.state_dict())

代码示例:(完整代码看我GitHub)

if torch.cuda.device_count() > 1:  # 判断GPU卡是否大于1
    net = AlexNet(num_classes=5, init_weights=True)
    net = nn.DataParallel(net.cuda(), device_ids=[0, 1])  # 包装model,并指定GPU

    epochs = 10
    save_path = './AlexNet.pth'
    best_acc = 0.0
    train_steps = len(train_loader)
    for epoch in range(epochs):
        # train
        net.train()
        running_loss = 0.0
        train_bar = tqdm(train_loader, file=sys.stdout)
        for step, data in enumerate(train_bar):
            images, labels = data
            optimizer.zero_grad()
            outputs = net(images.cuda()) # image,labels复制到cuda中
            loss = loss_function(outputs, labels.cuda())
            loss.backward()
            optimizer.step()
            
            # print statistics
            running_loss += loss.item()
            
            train_bar.desc = "train epoch[{}/{}] loss:{:.3f}".format(epoch + 1,
                                                                     epochs,
                                                                     loss)
        # validate
        net.eval()
        acc = 0.0  # accumulate accurate number / epoch
        with torch.no_grad():
            val_bar = tqdm(validate_loader, file=sys.stdout)
            for val_data in val_bar:
                val_images, val_labels = val_data
                outputs = net(val_images.cuda())
                predict_y = torch.max(outputs, dim=1)[1]
                acc += torch.eq(predict_y, val_labels.cuda()).sum().item()
        
        val_accurate = acc / val_num
        print('[epoch %d] train_loss: %.3f  val_accuracy: %.3f' %
              (epoch + 1, running_loss / train_steps, val_accurate))
        
        if val_accurate > best_acc:
            best_acc = val_accurate
            torch.save(net.state_dict(), save_path)
  • 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

2.torch.nn.parallel.DistributedDataParallel

DDP模式:该方法用命令行方式来运行代码
python -m torch.distributed.launch --nproc_per_node= n_gpus train.py
torch.distributed.launch: 用于控制当前节点哪个gpu运行,可以有argparse模块获取;
n_gpus:表示用几块gpu卡训练。

代码示例如下:

import json
from torch.utils.data import Dataset, DataLoader, DistributedSampler
import argparse
import torch
import torch.nn as nn
from torchvision import transforms, datasets

import sys
import os
import logging

logging.basicConfig(
    level=logging.WARN,
    stream=sys.stdout,
    format="%(asctime)s (%(module)s:%(lineno)d) %(levelname)s: %(message)s",
)


# 模型定义
class AlexNet(nn.Module):
    def __init__(self, num_classes=1000, init_weights=False):
        super(AlexNet, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 48, kernel_size=11, stride=4, padding=2),  # input[3, 224, 224]  output[48, 55, 55]
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),  # output[48, 27, 27]
            nn.Conv2d(48, 128, kernel_size=5, padding=2),  # output[128, 27, 27]
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),  # output[128, 13, 13]
            nn.Conv2d(128, 192, kernel_size=3, padding=1),  # output[192, 13, 13]
            nn.ReLU(inplace=True),
            nn.Conv2d(192, 192, kernel_size=3, padding=1),  # output[192, 13, 13]
            nn.ReLU(inplace=True),
            nn.Conv2d(192, 128, kernel_size=3, padding=1),  # output[128, 13, 13]
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),  # output[128, 6, 6]
        )
        self.classifier = nn.Sequential(
            nn.Dropout(p=0.5),
            nn.Linear(128 * 6 * 6, 2048),
            nn.ReLU(inplace=True),
            nn.Dropout(p=0.5),
            nn.Linear(2048, 2048),
            nn.ReLU(inplace=True),
            nn.Linear(2048, num_classes),
        )
        if init_weights:
            self._initialize_weights()
    
    def forward(self, x):
        x = self.features(x)
        x = torch.flatten(x, start_dim=1)
        x = self.classifier(x)
        return x
    
    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, 0, 0.01)
                nn.init.constant_(m.bias, 0)


# trian
def train(local_rank, train_dataset, eval_dataset, model, optimizer, num_epoch, log_step_interval, save_step_interval,
          eval_step_interval, save_path):
    """ 此处data_loader是map-style dataset """
    batch_size = 32
    nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, 8])  # number of workers
    
    model = nn.parallel.DistributedDataParallel(model.cuda(local_rank),
                                                device_ids=[local_rank])  # 模型拷贝,放入DistributedDataParallel
    
    train_sampler = DistributedSampler(train_dataset)
    train_data_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, sampler=train_sampler,
                                                    num_workers=nw)
    validate_loader = torch.utils.data.DataLoader(eval_dataset, batch_size=8)
    val_num = len(validate_dataset)
    train_steps = len(train_data_loader)
    for epoch_index in range(num_epoch):
        model.train()
        
        train_sampler.set_epoch(epoch_index)  # 为了让每张卡在每个周期中得到的数据是随机的
        
        for batch_index, (images, labels) in enumerate(train_data_loader):
            optimizer.zero_grad()
            running_loss = 0.0
            images = images.cuda(local_rank)  # 数据拷贝
            labels = labels.cuda(local_rank)  # 数据拷贝
            
            outputs = model(images)
            loss = nn.CrossEntropyLoss()(outputs, labels)
            loss.backward()
            optimizer.step()
            
            # print statistics
            running_loss += loss.item()
            
            # validate
        model.eval()
        acc = 0.0  # accumulate accurate number / epoch
        with torch.no_grad():
            for val_data in validate_loader:
                val_images, val_labels = val_data
                outputs = model(val_images)
                predict_y = torch.max(outputs, dim=1)[1]
                acc += torch.eq(predict_y, val_labels).sum().item()
        
        val_accurate = acc / val_num
        print('[epoch %d] train_loss: %.3f  val_accuracy: %.3f' %
              (epoch_index + 1, running_loss / train_steps, val_accurate))
        
        if val_accurate > best_acc:
            best_acc = val_accurate
            torch.save(model.modules().state_dict(), save_path)


# test
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--local_rank", help="local device id on current node",
                        type=int)
    args = parser.parse_args()
    
    if torch.cuda.is_available():
        logging.warning("Cuda is available!")
        if torch.cuda.device_count() > 1:
            logging.warning(f"Find {torch.cuda.device_count()} GPUs!")
        else:
            logging.warning("Too few GPU!")
    else:
        logging.warning("Cuda is not available! Exit!")
    
    n_gpus = 2
    torch.distributed.init_process_group("nccl", world_size=n_gpus, rank=args.local_rank)
    torch.cuda.set_device(args.local_rank)
    
    data_transform = {
        "train": transforms.Compose([transforms.RandomResizedCrop(224),
                                     transforms.RandomHorizontalFlip(),
                                     transforms.ToTensor(),
                                     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]),
        "val": transforms.Compose([transforms.Resize((224, 224)),
                                   transforms.ToTensor(),
                                   transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])}
    
    data_root = os.path.abspath(os.path.join(os.getcwd(), "../.."))  # get data root path
    image_path = os.path.join(data_root, "data_set", "flower_data")  # flower data set path
    assert os.path.exists(image_path), "{} path does not exist.".format(image_path)
    train_dataset = datasets.ImageFolder(root=os.path.join(image_path, "train"),
                                         transform=data_transform["train"])
    train_num = len(train_dataset)
    
    # {'daisy':0, 'dandelion':1, 'roses':2, 'sunflower':3, 'tulips':4}
    flower_list = train_dataset.class_to_idx
    cla_dict = dict((val, key) for key, val in flower_list.items())
    # write dict into json file
    json_str = json.dumps(cla_dict, indent=4)
    with open('class_indices.json', 'w') as json_file:
        json_file.write(json_str)
    
    validate_dataset = datasets.ImageFolder(root=os.path.join(image_path, "val"),
                                            transform=data_transform["val"])
    
    model = AlexNet()
    
    print("模型总参数:", sum(p.numel() for p in model.parameters()))
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    train(args.local_rank, train_dataset, validate_dataset, model,
          optimizer, num_epoch=10, log_step_interval=20, save_step_interval=500, eval_step_interval=300,
          save_path="./")

  • 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
  • 67
  • 68
  • 69
  • 70
  • 71
  • 72
  • 73
  • 74
  • 75
  • 76
  • 77
  • 78
  • 79
  • 80
  • 81
  • 82
  • 83
  • 84
  • 85
  • 86
  • 87
  • 88
  • 89
  • 90
  • 91
  • 92
  • 93
  • 94
  • 95
  • 96
  • 97
  • 98
  • 99
  • 100
  • 101
  • 102
  • 103
  • 104
  • 105
  • 106
  • 107
  • 108
  • 109
  • 110
  • 111
  • 112
  • 113
  • 114
  • 115
  • 116
  • 117
  • 118
  • 119
  • 120
  • 121
  • 122
  • 123
  • 124
  • 125
  • 126
  • 127
  • 128
  • 129
  • 130
  • 131
  • 132
  • 133
  • 134
  • 135
  • 136
  • 137
  • 138
  • 139
  • 140
  • 141
  • 142
  • 143
  • 144
  • 145
  • 146
  • 147
  • 148
  • 149
  • 150
  • 151
  • 152
  • 153
  • 154
  • 155
  • 156
  • 157
  • 158
  • 159
  • 160
  • 161
  • 162
  • 163
  • 164
  • 165
  • 166
  • 167
  • 168
  • 169
  • 170
  • 171
  • 172
  • 173
  • 174
  • 175

执行命令:
python -m torch.distributed.launch --nproc_per_node=2 trian.py


本文内容由网友自发贡献,转载请注明出处:https://www.wpsshop.cn/w/煮酒与君饮/article/detail/932972
推荐阅读
相关标签
  

闽ICP备14008679号