16、Pytorch Lightning入门

2022-10-21 10:45:45

资源

官方手册
GitHub地址
GItHub案例:Pytorch-Lightning-Template项目
pytorch也是有缺陷的,例如要用半精度训练、BatchNorm参数同步、单机多卡训练,则要安排一下Apex。而pl则不同,这些全部都安排,而且只要设置一下参数就可以了。另外,还有一个特色,就是你的超参数全部保存到模型中,如果你要调巨多参数,那就不需要再对每个训练的模型进行参数标记了,而且恢复模型时可以直接恢复超参数,可以大大减小代码量和工作量

基础案例实现ResNet

import torchfrom torchimport nnfrom torchimport optimfrom torchvisionimport datasets, transformsfrom torch.utils.dataimport random_split, DataLoaderimport pytorch_lightningas plfrom pytorch_lightning.metrics.functionalimport accuracyclassResNet(pl.LightningModule):def__init__(self):super().__init__()
        self.l1= nn.Linear(28*28,64)
        self.l2= nn.Linear(64,64)
        self.l3= nn.Linear(64,10)
        self.do= nn.Dropout(0.1)
        self.loss= nn.CrossEntropyLoss()defforward(self, x):# nn.ReLU()
        h1= nn.functional.relu(self.l1(x))
        h2= nn.functional.relu(self.l2(h1))
        do= self.do(h2+ h1)
        logits= self.l3(do)return logitsdefconfigure_optimizers(self):
        optimiser= optim.SGD(self.parameters(), lr=1e-2)return optimiserdeftraining_step(self, batch, batch_idx):
        x, y= batch
        b= x.size(0)
        x= x.view(b,-1)
        logits= self(x)
        J= self.loss(logits, y)

        acc= accuracy(logits, y)
        pbar={"train_acc": acc}return{"loss": J,"progress_bar": pbar}defvalidation_step(self, batch, batch_idx):
        results= self.training_step(batch, batch_idx)
        results['progress_bar']['val_acc']= results['progress_bar']['train_acc']del results['progress_bar']['train_acc']return resultsdefvalidation_epoch_end(self, val_step_outputs):
        avg_val_loss= torch.tensor([x['loss']for xin val_step_outputs]).mean()
        avg_val_acc= torch.tensor([x['progress_bar']['val_acc']for xin val_step_outputs]).mean()

        pbar={"avg_val_acc": avg_val_acc}return{'val_loss': avg_val_loss,'progress_bar': pbar}defprepare_data(self):
        datasets.MNIST('data', train=True, download=True, transform=transforms.ToTensor())defsetup(self, stage:str):
        train_data= datasets.MNIST('data', train=True, download=False, transform=transforms.ToTensor())
        self.tra, self.val= random_split(train_data,[55000,5000])deftrain_dataloader(self):
        train_loader= DataLoader(self.tra, batch_size=32, num_workers=48)return train_loaderdefval_dataloader(self):
        val_loader= DataLoader(self.val, batch_size=32, num_workers=48)return val_loaderif __name__=='__main__':
    model= ResNet()# Lightning 会自动保存最近训练的epoch的模型到当前的工作空间(or.getcwd()),也可以在定义Trainer的时候指定
    trainer= pl.Trainer(progress_bar_refresh_rate=20, max_epochs=1, gpus=1, default_root_dir="./root")
    trainer.fit(model)

数据集

数据集有两种实现方法,直接调用第三方公开数据集(如:MNIST等数据集)和 自定义数据集(继承torch.utils.data.dataset.Dataset)

使用现有的公开数据集

defprepare_data(self):
        datasets.MNIST('data', train=True, download=True, transform=transforms.ToTensor())defsetup(self, stage:str):
        train_data= datasets.MNIST('data', train=True, download=False, transform=transforms.ToTensor())
        self.tra, self.val, self.test= random_split(train_data,[50000,5000,5000])deftrain_dataloader(self):
        train_loader= DataLoader(self.tra, batch_size=32, num_workers=48)return train_loaderdefval_dataloader(self):
        val_loader= DataLoader(self.val, batch_size=32, num_workers=48)return val_loaderdeftest_dataloader(self):
        test_loader= DataLoader(self.test, batch_size=32, num_workers=48)return test_loader

自定义数据集

Dataset

import sysimport pathlibimport torchfrom torch.utils.dataimport Datasetfrom utilsimport sort_batch_by_len, source2ids

abs_path= pathlib.Path(__file__).parent.absolute()
sys.path.append(sys.path.append(abs_path))classSampleDataset(Dataset):"""
    The class represents a sample set for training.
    """def__init__(self, data_pairs, vocab):
        self.src_texts=[data_pair[0]for data_pairin data_pairs]
        self.tgt_texts=[data_pair[1]for data_pairin data_pairs]
        self.vocab= vocab
        self._len=len(data_pairs)# Keep track of how many data points.def__len__(self):return self._lendef__getitem__(self, index):# print("\nself.src_texts[{0}] = {1}".format(index, self.src_texts[index]))
        src_ids, oovs= source2ids(self.src_texts[index], self.vocab)# 将当前文本self.src_texts[index]转为ids,oovs为超出词典范围的词汇文本
        item={'x':[self.vocab.SOS]+ src_ids+[self.vocab.EOS],'y':[self.vocab.SOS]+[self.vocab[i]for iin self.tgt_texts[index]]+[self.vocab.EOS],'x_len':len(self.src_texts[index]),'y_len':len(self.tgt_texts[index]),'oovs': oovs,'len_oovs':len(oovs)}return item

DataLoader

from torch.utils.dataimport DataLoader, random_splitimport pytorch_lightningas plclassMyDataModule(pl.LightningDataModule):def__init__(self):super().__init__()defprepare_data(self):# 在该函数里一般实现数据集的下载等,只有cuda:0 会执行该函数# download, split, etc...# only called on 1 GPU/TPU in distributedpassdefsetup(self, stage):# make assignments here (val/train/test split)# called on every process in DDP# 实现数据集的定义,每张GPU都会执行该函数, stage 用于标记是用于什么阶段if stage=='fit'or stageisNone:
            self.train_dataset= MyDataset(self.train_file_path, self.train_file_num, transform=None)
            self.val_dataset= MyDataset(self.val_file_path, self.val_file_num, transform=None)if stage=='test'or stageisNone:
            self.test_dataset= MyDataset(self.test_file_path, self.test_file_num, transform=None)deftrain_dataloader(self):return DataLoader(self.train_dataset, batch_size=self.batch_size, shuffle=False, num_workers=0)defval_dataloader(self):return DataLoader(self.val_dataset, batch_size=self.batch_size, shuffle=False)deftest_dataloader(self):return DataLoader(self.test_dataset, batch_size=1, shuffle=True)

pytorch-lightning流程

初始化 definit(self) -->训练training_step(self, batch, batch_idx) --> 校验validation_step(self, batch, batch_idx) --> 测试 test_step(self, batch, batch_idx). 就完事了

当然,除了这三个主要的,还有一些其他的函数,为了方便我们实现其他的一些功能,因此更为完整的流程是在training_step 、validation_step、test_step 后面都紧跟着其相应的 training_step_end(self,batch_parts)和training_epoch_end(self, training_step_outputs) 函数,当然,对于校验和测试,都有相应的_step_end和_epoch_end函数。**

** *_step_end – 即每一个 * 步完成后调用 **

** *_epoch_end – 即每一个 * 的epoch 完成之后会自动调用 **

deftraining_step(self, batch, batch_idx):
    x, y= batch
    y_hat= self.model(x)
    loss= F.cross_entropy(y_hat, y)
    pred=...return{'loss': loss,'pred': pred}deftraining_step_end(self, batch_parts):'''
    当gpus=0 or 1时,这里的batch_parts即为traing_step的返回值(已验证)
    当gpus>1时,这里的batch_parts为list,list中每个为training_step返回值,list[i]为i号gpu的返回值(这里未验证)
    '''
    gpu_0_prediction= batch_parts[0]['pred']
    gpu_1_prediction= batch_parts[1]['pred']# do something with both outputsreturn(batch_parts[0]['loss']+ batch_parts[1]['loss'])/2deftraining_epoch_end(self, training_step_outputs):'''
    当gpu=0 or 1时,training_step_outputs为list,长度为steps的数量(不包括validation的步数,当你训练时,你会发现返回list<训练时的steps数,这是因为训练时显示的steps数据还包括了validation的,若将limit_val_batches=0.,即关闭validation,则显示的steps会与training_step_outputs的长度相同)。list中的每个值为字典类型,字典中会存有`training_step_end()`返回的键值,键名为`training_step()`函数返回的变量名,另外还有该值是在哪台设备上(哪张GPU上),例如{device='cuda:0'}
    '''for outin training_step_outputs:# do something with preds

Train

训练主要是重写def training_setp(self, batch, batch_idx)函数,并返回要反向传播的loss即可,其中batch 即为从 train_dataloader 采样的一个batch的数据,batch_idx即为目前batch的索引

deftraining_setp(self, batch, batch_idx):
    image, label= batch
    pred= self.forward(iamge)
    loss=...# 一定要返回lossreturn loss

Validation

  • 每训练n个epochs 校验一次
    默认为每1个epoch校验一次,即自动调用validation_step()函数:check_val_every_n_epoch
trainer= Trainer(check_val_every_n_epoch=1)
  • 单个epoch内校验频率
    当一个epoch 比较大时,就需要在单个epoch 内进行多次校验,这时就需要对校验的调动频率进行修改, 传入val_check_interval的参数为float型时表示百分比,为int时表示batch
# 每训练单个epoch的 25% 调用校验函数一次,注意:要传入float型数
trainer= Trainer(val_check_interval=0.25)# 当然也可以是单个epoch训练完多少个batch后调用一次校验函数,但是一定是传入int型
trainer= Trainer(val_check_interval=100)# 每训练100个batch校验一次

校验和训练是一样的,重写def validation_step(self, batch, batch_idx)函数

defvalidation_step(self, batch, batch_idx):
    results= self.training_step(batch, batch_idx)
    results['progress_bar']['val_acc']= results['progress_bar']['train_acc']del results['progress_bar']['train_acc']return results

test

在pytoch_lightning框架中,test 在训练过程中是不调用的,也就是说是不相关,在训练过程中只进行training和validation,因此如果需要在训练过中保存validation的一些信息,就要放到validation中。测试是在训练完成之后的

# 获取恢复了权重和超参数等的模型
model= MODEL.load_from_checkpoint(checkpoint_path='my_model_path/heiheihei.ckpt')# 修改测试时需要的参数,例如预测的步数等
model.pred_step=1000# 定义trainer, 其中limit_test_batches表示取测试集中的0.05的数据来做测试
trainer= pl.Trainer(gpus=1, precision=16, limit_test_batches=0.05)# 测试,自动调用test_step(), 其中dm为数据集
trainer.test(model=dck, datamodule=dm)

模型、Trainer的保存与恢复

Lightning 会自动保存最近训练的epoch的模型到当前的工作空间(or.getcwd()),也可以在定义Trainer的时候指定

trainer= Trainer(default_root_dir='/your/path/to/save/checkpoints')

也可以关闭自动保存模型

trainer= Trainer(checkpoint_callback=False)

利用ModelCheckpoint (callbacks)保存模型

所有参数均为optional

ModelCheckpoint(
    dirpath=None,
    filename=None,
    monitor=None,
    verbose=False,
    save_last=None,
    save_top_k=1,
    save_weights_only=False,
    mode="min",
    auto_insert_metric_name=True,
    every_n_train_steps=None,
    train_time_interval=None,
    every_n_epochs=None,
    save_on_train_epoch_end=None,
    every_n_val_epochs=None)

在这里插入图片描述
自动保存下,也可以自定义要监控的量来保存模型

  • 计算需要监控的量,例如校验误差:loss
  • 使用log()函数标记该要监控的量(直接在training_step、validation_step中添加)
  • 初始化ModelCheckpoint回调,并设置要监控的量
  • 将其传回到Trainer中
from pytorch_lightningimport Trainer, LightningDataModule, LightningModule, Callback, seed_everythingfrom pytorch_lightning.callbacksimport ModelCheckpointfrom pytorch_lightning.loggersimport TensorBoardLogger# Transformer LightningModuleclassGLUETransformer(LightningModule):...deftraining_step(self, batch, batch_idx):# 1. 计算loss
        outputs= self(**batch)
        train_loss= outputs[0]# 2. 使用log()函数标记该要监控的量,名字叫'val_loss'
        self.log('ltrain_lossoss', train_loss)return train_loss...defvalidation_step(self, batch, batch_idx, dataloader_idx=0):
        outputs= self(**batch)# 1. 计算需要监控的量
        val_loss, logits= outputs[:2]# 2. 使用log()函数标记该要监控的量,名字叫'val_loss'
        self.log('val_loss', val_loss)if self.hparams.num_labels>=1:
            preds= torch.argmax(logits, axis=1)elif self.hparams.num_labels==1:
            preds= logits.squeeze()

        labels= batch["labels"]return{"loss": val_loss,"preds": preds,"labels": labels}...# Training & Testif __name__=="__main__":
    seed_everything(42)# 定义数据集
    data_module= GLUEDataModule(model_name_or_path=r"D:\Pretrained_Model\albert-base-v2", task_name="cola")# 定义模型
    model= GLUETransformer(model_name_or_path=r"D:\Pretrained_Model\albert-base-v2", num_labels=data_module.num_labels, eval_splits=data_module.eval_splits, task_name=data_module.task_name)# 初始化`ModelCheckpoint`回调,并设置要监控的量
    checkpoint_callback= ModelCheckpoint(
        dirpath='saved_module',
        filename='sample-cola-{epoch:02d}-{val_loss:.2f}',
        monitor='val_loss')# 定义trainer
    trainer= Trainer(max_epochs=20, gpus=AVAIL_GPUS, check_val_every_n_epoch=1, callbacks=[checkpoint_callback])# 默认为每1个epoch校验一次,即自动调用validation_step()函数;将 checkpoint_callback 放到Trainer的callback 的list中# 开始训练
    trainer.fit(model=model, datamodule=data_module)

在上面的 filename 参数中,定义了模型文件的保存格式,然后通过自动调用format_checkpoint_name 函数给其中的变量赋值的,返回 string 类型,文件名

>>> tmpdir= os.path.dirname(__file__)>>> ckpt= ModelCheckpoint(dirpath=tmpdir, filename='{epoch}')>>> os.path.basename(ckpt.format_checkpoint_name(0,1, metrics={}))'epoch=0.ckpt'>>> ckpt= ModelCheckpoint(dirpath=tmpdir, filename='{epoch:03d}')>>> os.path.basename(ckpt.format_checkpoint_name(5,2, metrics={}))'epoch=005.ckpt'>>> ckpt= ModelCheckpoint(dirpath=tmpdir, filename='{epoch}-{val_loss:.2f}')>>> os.path.basename(ckpt.format_checkpoint_name(2,3, metrics=dict(val_loss=0.123456)))'epoch=2-val_loss=0.12.ckpt'>>> ckpt= ModelCheckpoint(dirpath=tmpdir
  • 作者:C--G
  • 原文链接:https://blog.csdn.net/weixin_50973728/article/details/126227752
    更新时间:2022-10-21 10:45:45