当前位置:   article > 正文

pytorch单机多GPU训练方式(微调bert)_torch多gpu微调bert

torch多gpu微调bert

本来是用Dataparallel来微调,但是发现gpu0上负载太大了,后来上面一看,pytorch官网推荐使用DistributedDataParallel,这个方法可以用于多机多卡或者单机多卡,速度和各方面都比Dataparallel要好很多。
我主要看了下面两个博客:
【分布式训练】单机多卡的正确打开方式(三):PyTorch
pytorch(分布式)数据并行个人实践总结——DataParallel/DistributedDataParallel

这里说一下Dataparallel和DistributedDataparallel在使用时的区别。比如有三张卡(也就是3张gpu),Batch=30。

  1. Dataparallel只开一个进程,DistributedDataparallel有几张卡开几个进程,不过需要在执行程序时手动设置。
  2. Dataparallel返回的loss是3张卡的loss,所以需要mean一下,这样才是整个batch的loss,也就是Dataparallel会自动将输入的一个batch的数据切分成3份,每张卡分别计算一部分,也就是10份。
    DistributedDataparallel是不会自动切分数据的,所以在Dataloader中,需要将sampler设置为DistributedSampler,具体见下面的代码。DistributedDataparallel是将整个数据集切分为3份,然后每一个进程使用其中的一份,但是每个进程的batch还是30。
  3. 在保存模型时,Dataparallel和DistributedDataparallel都需要使用model.module.state_dict来保存模型,而不是model.state_dict,不过DistributedDataparallel还需要选择一个进程的模型保存,具体见代码。
  4. 执行DistributedDataparallel的程序是要用命令行执行,python -m torch.distributed.launch --nproc_per_node=3 main.py, 其中 nproc_per_node指定了要用为每一张卡开一个进程。
  5. Dataparallel只开一个进程,默认是0号GPU为master,也就是模型会放在0号GPU上,然后其他GPU上是模型的浅拷贝,步骤2中得到的数据会分别送到每个GPU上独立并行的进行forward过程,然后主GPU会收集loss,loss平均后分发给每个GPU进行反向传播(但此时并不更新参数,仅仅是反向传播计算梯度),然后各GPU将梯度汇总到主GPU,进行梯度下降,并更新模型参数,然后将参数广播到其余GPU。可以看到,在上述过程中,主GPU做了很多事,所以在这种并行方式下,主GPU的显卡占用率要比其他gpu高很多。
  6. DistributedDataparallel为每个GPU开一个进程,没有主GPU,每个GPU执行相同的任务,每个进程具有自己的 optimizer ,并独立完成所有的优化步骤,进程内与一般的训练无异。在各进程梯度计算完成之后,各进程需要将梯度进行汇总平均,然后再由 rank=0 的进程,将其 broadcast 到所有进程。之后,各进程用该梯度来独立的更新参数。相较于 DataparallelDistributedDataparallel传输的数据量更少,因此速度更快,效率更高。
import random
import numpy as np
from transformers import AlbertTokenizer, AlbertForMaskedLM, AdamW, get_linear_schedule_with_warmup, \
    DataCollatorForLanguageModeling
from torch.nn.parallel import DistributedDataParallel
from torch.utils.data import Dataset, DataLoader, RandomSampler
from torch.utils.data.distributed import DistributedSampler
import torch
import time
from utils import format_time
from utils import read_data

torch.distributed.init_process_group(backend='nccl')

tokenizer = AlbertTokenizer.from_pretrained('albert-base-v2', cache_dir='../language_model/albert')
model = AlbertForMaskedLM.from_pretrained('../language_model/albert')

seed_val = 42
random.seed(seed_val)
np.random.seed(seed_val)
torch.manual_seed(seed_val)
torch.cuda.manual_seed_all(seed_val)


def my_collate_fn(batch):
    input_ids = []
    attention_masks = []
    for sent in batch:
        encoded_dict = tokenizer.encode_plus(
            sent[0],  # 输入文本
            add_special_tokens=True,  # 添加 '[CLS]' 和 '[SEP]'
            max_length=128,  # 填充 & 截断长度
            truncation=True,
            pad_to_max_length=True,
            return_attention_mask=True,  # 返回 attn. masks.
        )
        input_ids.append(torch.tensor(encoded_dict['input_ids']))

        attention_masks.append(torch.tensor(encoded_dict['attention_mask']))
    input_ids = torch.stack(input_ids, dim=0)
    attention_masks = torch.stack(attention_masks, dim=0)
    return [input_ids, attention_masks]


class myDataset(Dataset):

    def __init__(self, data):
        self.data = data

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

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


def make_data(batch_size):
    data = read_data('en')
    dataset = myDataset(data)

    data_dataloader = DataLoader(
        dataset,
        sampler=DistributedSampler(dataset),  # 这里也是分布式的sampler
        batch_size=batch_size,
        collate_fn=my_collate_fn,
        num_workers=4,
    )

    return data_dataloader

'''
这个方法也是bert自带的实现了随机mask 15%单词任务的方法。
由于我是要在领域内的数据集上微调bert,所以还是用language model的训练方法。
'''
datacollecter = DataCollatorForLanguageModeling(tokenizer)

def trian(model, train_loader, optimizer, scheduler, local_rank):
    for epoch_i in range(0, epochs):

        # ========================================
        #               Training
        # ========================================

        print("")
        print('======== Epoch {:} / {:} ========'.format(epoch_i + 1, epochs))
        print('Training...')

        total_train_loss = 0
        model.train()
        t0 = time.time()

        for step, batch in enumerate(train_loader):

            if step % 50 == 0 and not step == 0:
                elapsed = format_time(time.time() - t0)
                print('  Batch {:>5,}  of  {:>5,}.    Elapsed: {:}.'.format(step, len(train_loader), elapsed))

            b_input_ids, b_attention_mask = batch[0], batch[1]
            output = datacollecter(b_input_ids)
            b_input_ids, b_labels = output["input_ids"], output["labels"]
            b_input_ids = b_input_ids.to(device)
            b_attention_mask = b_attention_mask.to(device)

            model.zero_grad()
            loss, logits = model(b_input_ids, attention_mask=b_attention_mask, labels=b_labels)

            total_train_loss += loss.item()

            loss.backward()
            optimizer.step()
            scheduler.step()

        avg_train_loss = total_train_loss / len(train_loader)

        training_time = format_time(time.time() - t0)

        print("")
        print("  Average training loss: {0:.2f}".format(avg_train_loss))
        print("  Training epcoh took: {:}".format(training_time))

    ## 选择一个进程保存
    if local_rank == 0:
        model.module.save_pretrained(save_path)  # save_pretrained是bert自带的保存微调模型的方法
        print('Saving model in %s.' % save_path)

    print("")
    print("Training complete!")
    print("Total training took {:} (h:mm:ss)".format(format_time(time.time() - total_t0)))


epochs = 2
batch_size = 128
save_path = './model/'


if __name__ == "__main__":

    # 为每个进程配置GPU
    local_rank = torch.distributed.get_rank()
    torch.cuda.set_device(local_rank)
    device = torch.device("cuda", local_rank)

    # 要先将model放到gpu上
    model = model.to(device)
    model = DistributedDataParallel(model, find_unused_parameters=True, device_ids=[local_rank],
                                    output_device=local_rank)

    training_stats = []
    total_t0 = time.time()

    train_loader = make_data(batch_size)
    total_steps = len(train_loader) * epochs

    optimizer = AdamW(model.parameters(),
                      lr=2e-5,  # args.learning_rate - default is 5e-5
                      eps=1e-8  # args.adam_epsilon  - default is 1e-8
                      )

    scheduler = get_linear_schedule_with_warmup(optimizer,
                                                num_warmup_steps=1000,
                                                num_training_steps=total_steps)

    trian(model, train_loader, optimizer, scheduler, local_rank)


  • 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
声明:本文内容由网友自发贡献,不代表【wpsshop博客】立场,版权归原作者所有,本站不承担相应法律责任。如您发现有侵权的内容,请联系我们。转载请注明出处:https://www.wpsshop.cn/w/我家自动化/article/detail/708056
推荐阅读
相关标签
  

闽ICP备14008679号