Pytorch-lightning的使用案例

小咪咪 2022-10-12 00:54 314阅读 0赞

点击上方“**机器学习与生成对抗网络**”，关注星标

获取有趣、好玩的前沿干货！

作者丨Caliber@知乎  编辑丨极市平台

来源丨https://zhuanlan.zhihu.com/p/370185203

Pytorch-lightning(以下简称pl)可以非常简洁得构建深度学习代码。但是其实大部分人用不到很多复杂得功能。而pl有时候包装得过于深了，用的时候稍微有一些不灵活。通常来说，在你的模型搭建好之后，大部分的功能都会被封装在一个叫trainer的类里面。一些比较麻烦但是需要的功能通常如下：

1.  保存checkpoints
2.  输出log信息
3.  resume training 即重载训练，我们希望可以接着上一次的epoch继续训练
4.  记录模型训练的过程(通常使用tensorboard)
5.  设置seed，即保证训练过程可以复制

好在这些功能在pl中都已经实现。

由于doc上的很多解释并不是很清楚，而且网上例子也不是特别多。下面分享一点我自己的使用心得。

首先关于设置全局的种子：

from pytorch_lightning import seed_everything
    # Set seedseed = 42seed_everything(seed)

只需要import如上的seed\_everything函数即可。它应该和如下的函数是等价的：

def seed_all(seed_value):    random.seed(seed_value) # Python    np.random.seed(seed_value) # cpu vars    torch.manual_seed(seed_value) # cpu vars
        if torch.cuda.is_available():         print ('CUDA is available')        torch.cuda.manual_seed(seed_value)        torch.cuda.manual_seed_all(seed_value) # gpu vars        torch.backends.cudnn.deterministic = True  #needed        torch.backends.cudnn.benchmark = False
    seed=42seed_all(seed)

但经过我的测试，好像pl的seed\_everything函数应该更全一点。

下面通过一个具体的例子来说明一些使用方法：

先下载、导入必要的包和下载数据集：

!pip install pytorch-lightning!wget https://download.pytorch.org/tutorial/hymenoptera_data.zip!unzip -q hymenoptera_data.zip!rm hymenoptera_data.zip
    import pytorch_lightning as plimport osimport numpy as np import randomimport matplotlib.pyplot as plt
    import torchimport torch.nn.functional as Fimport torchvisionimport torchvision.transforms as transforms

以下代码种加入!的代码是在terminal中运行的。在google colab中运行linux命令需要在之前加!

如果是使用google colab，由于它创建的是一个虚拟机，不能及时保存，所以如果需要保存，挂载自己google云盘也是有必要的。使用如下的代码：

from google.colab import drivedrive.mount('./content/drive')
    import osos.chdir("/content/drive/My Drive/")

先如下定义如下的LightningModule和main函数。

class CoolSystem(pl.LightningModule):
    
    
        def __init__(self, hparams):
            super(CoolSystem, self).__init__()
    
    
            self.params = hparams
    
    
            self.data_dir = self.params.data_dir
            self.num_classes = self.params.num_classes 
    
    
            ########## define the model ########## 
            arch = torchvision.models.resnet18(pretrained=True)
            num_ftrs = arch.fc.in_features
    
    
            modules = list(arch.children())[:-1] # ResNet18 has 10 children
            self.backbone = torch.nn.Sequential(*modules) # [bs, 512, 1, 1]
            self.final = torch.nn.Sequential(
                   torch.nn.Linear(num_ftrs, 128),
                   torch.nn.ReLU(inplace=True),
                   torch.nn.Linear(128, self.num_classes),
                   torch.nn.Softmax(dim=1))
    
    
        def forward(self, x):
            x = self.backbone(x)
            x = x.reshape(x.size(0), -1)
            x = self.final(x)
    
    
            return x
    
    
        def configure_optimizers(self):
            # REQUIRED
            optimizer = torch.optim.SGD([
                    {'params': self.backbone.parameters()},
                    {'params': self.final.parameters(), 'lr': 1e-2}
                ], lr=1e-3, momentum=0.9)
    
    
            exp_lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)
    
    
            return [optimizer], [exp_lr_scheduler]
    
    
        def training_step(self, batch, batch_idx):
            # REQUIRED
            x, y = batch
            y_hat = self.forward(x)
    
    
            loss = F.cross_entropy(y_hat, y)
    
    
            _, preds = torch.max(y_hat, dim=1)
            acc = torch.sum(preds == y.data) / (y.shape[0] * 1.0)
    
    
            self.log('train_loss', loss)
            self.log('train_acc', acc)
    
    
            return {'loss': loss, 'train_acc': acc}
    
    
    
    
        def validation_step(self, batch, batch_idx):
            # OPTIONAL
            x, y = batch
            y_hat = self.forward(x)
            loss = F.cross_entropy(y_hat, y)
            _, preds = torch.max(y_hat, 1)
            acc = torch.sum(preds == y.data) / (y.shape[0] * 1.0)
    
    
            self.log('val_loss', loss)
            self.log('val_acc', acc)
    
    
            return {'val_loss': loss, 'val_acc': acc}
    
    
    
    
        def test_step(self, batch, batch_idx):
            # OPTIONAL
            x, y = batch
            y_hat = self.forward(x)
            loss = F.cross_entropy(y_hat, y)
            _, preds = torch.max(y_hat, 1)
            acc = torch.sum(preds == y.data) / (y.shape[0] * 1.0)
    
    
            return {'test_loss': loss, 'test_acc': acc}
    
    
    
    
        def train_dataloader(self):
            # REQUIRED
    
    
            transform = transforms.Compose([
                                    transforms.RandomResizedCrop(224),
                                    transforms.RandomHorizontalFlip(),
                                    transforms.ToTensor(),
                                    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
                                    ])
    
    
            train_set = torchvision.datasets.ImageFolder(os.path.join(self.data_dir, 'train'), transform)
            train_loader = torch.utils.data.DataLoader(train_set, batch_size=32, shuffle=True, num_workers=4)
    
    
            return train_loader
    
    
        def val_dataloader(self):
          transform = transforms.Compose([
                                    transforms.Resize(256),
                                    transforms.CenterCrop(224),
                                    transforms.ToTensor(),
                                    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
                                    ])
    
    
          val_set = torchvision.datasets.ImageFolder(os.path.join(self.data_dir, 'val'), transform)
          val_loader = torch.utils.data.DataLoader(val_set, batch_size=32, shuffle=True, num_workers=4)
    
    
          return val_loader
    
    
        def test_dataloader(self):
          transform = transforms.Compose([
                                  transforms.Resize(256),
                                  transforms.CenterCrop(224),
                                  transforms.ToTensor(),
                                  transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
                                  ])
    
    
          val_set = torchvision.datasets.ImageFolder(os.path.join(self.data_dir, 'val'), transform)
          val_loader = torch.utils.data.DataLoader(val_set, batch_size=8, shuffle=True, num_workers=4)
    
    
          return val_loader
    
    
    
    
    
    
    
    
    
    
    def main(hparams):
      model = CoolSystem(hparams)
    
    
    
    
      trainer = pl.Trainer(
          max_epochs=hparams.epochs,
          gpus=1,
          accelerator='dp'
      )  
    
    
      trainer.fit(model)

下面是run的部分：

from argparse import Namespace
    args = {    'num_classes': 2,    'epochs': 5,    'data_dir': "/content/hymenoptera_data",}
    hyperparams = Namespace(**args)
    
    if __name__ == '__main__':    main(hyperparams)

如果希望重载训练的话，可以按如下方式：

# resume training
    RESUME = True
    if RESUME:    resume_checkpoint_dir = './lightning_logs/version_0/checkpoints/'    checkpoint_path = os.listdir(resume_checkpoint_dir)[0]    resume_checkpoint_path = resume_checkpoint_dir + checkpoint_path
    
        args = {    'num_classes': 2,    'data_dir': "/content/hymenoptera_data"}
        hparams = Namespace(**args)
        model = CoolSystem(hparams)
    
        trainer = pl.Trainer(gpus=1,                 max_epochs=10,                             accelerator='dp',                resume_from_checkpoint = resume_checkpoint_path)
        trainer.fit(model)

如果我们想要从checkpoint加载模型，并进行使用可以按如下操作来：

import matplotlib.pyplot as pltimport numpy as np
    # functions to show an imagedef imshow(inp):    inp = inp.numpy().transpose((1, 2, 0))    mean = np.array([0.485, 0.456, 0.406])    std = np.array([0.229, 0.224, 0.225])    inp = std * inp + mean    inp = np.clip(inp, 0, 1)    plt.imshow(inp)    plt.show()
    classes = ['ants', 'bees']
    checkpoint_dir = 'lightning_logs/version_1/checkpoints/'checkpoint_path = checkpoint_dir + os.listdir(checkpoint_dir)[0]
    checkpoint = torch.load(checkpoint_path)model_infer = CoolSystem(hparams)model_infer.load_state_dict(checkpoint['state_dict'])
    try_dataloader = model_infer.test_dataloader()
    inputs, labels = next(iter(try_dataloader))
    # print images and ground truthimshow(torchvision.utils.make_grid(inputs))print('GroundTruth: ', ' '.join('%5s' % classes[labels[j]] for j in range(8)))
    # inferenceoutputs = model_infer(inputs)
    _, preds = torch.max(outputs, dim=1)# print (preds)print (torch.sum(preds == labels.data) / (labels.shape[0] * 1.0))
    print('Predicted: ', ' '.join('%5s' % classes[preds[j]] for j in range(8)))

![69ed060d47290538fb78e7432006d537.png][]

预测结果如上。

如果希望检测训练过程（第一部分+重载训练的部分），如下：

# tensorboard
    %load_ext tensorboard%tensorboard --logdir = ./lightning_logs

![56cfcfa52fc9099b51d84ebeda64c4af.png][]

训练过程在tensorboard里面记录，version0是第一次的训练，version1是重载后的结果。

完整的code在这里.

https://colab.research.google.com/gist/calibertytz/a9de31175ce15f384dead94c2a9fad4d/pl\_tutorials\_1.ipynb

**猜您喜欢：**

**[CVPR 2021 | GAN的说话人驱动、3D人脸论文汇总][CVPR 2021 _ GAN_3D]**

**[CVPR 2021 | 图像转换 今如何？几篇GAN论文][CVPR 2021 _ _ _GAN]**

[**【CVPR 2021】通过GAN提升人脸识别的遗留难题**][CVPR 2021_GAN]

[**CVPR 2021生成对抗网络GAN部分论文汇总**][CVPR 2021_GAN 1]

[**经典GAN不得不读：StyleGAN**][GAN_StyleGAN]

[**最新最全20篇！基于 StyleGAN 改进或应用相关论文**][20_ StyleGAN]

**[超100篇！CVPR 2020最全GAN论文梳理汇总！][100_CVPR 2020_GAN]**

[**附下载 | 《Python进阶》中文版**][_ _Python]

[**附下载 | 经典《Think Python》中文版**][_ _Think Python]

[**附下载 | 《Pytorch模型训练实用教程》**][_ _Pytorch]

[**附下载 | 最新2020李沐《动手学深度学习》**][_ _2020]

**[附下载 | 《可解释的机器学习》中文版][Link 1]**

[**附下载 |《TensorFlow 2.0 深度学习算法实战》**][_TensorFlow 2.0]

**[附下载 | 超100篇！CVPR 2020最全GAN论文梳理汇总！][100_CVPR 2020_GAN]**

**[附下载 |《计算机视觉中的数学方法》分享][Link 2]**

![d9ab7c139edeea33a3b5c90e0ebc5b9d.png][]

[69ed060d47290538fb78e7432006d537.png]: /images/20221005/b8277d9c7716492994a097f74ddee943.png
[56cfcfa52fc9099b51d84ebeda64c4af.png]: /images/20221005/70cb2ca73a454e83bdee0ec0dbeab919.png
[CVPR 2021 _ GAN_3D]: http://mp.weixin.qq.com/s?__biz=MzU5MTgzNzE0MA%3D%3D&chksm=fe2a41fac95dc8ec752351699318742fa502dfc945ce072567bb4bf0c59ce83049dc58c1a57d&idx=1&mid=2247492417&scene=21&sn=292be5e5fab52c5df5c94508b898da5a#wechat_redirect
[CVPR 2021 _ _ _GAN]: http://mp.weixin.qq.com/s?__biz=MzU5MTgzNzE0MA%3D%3D&chksm=fe2a4278c95dcb6ebd0ba346a11dc195a0ac4f7830f3a7cf69fcffc7057fa1bf3eaeb9386300&idx=1&mid=2247491779&scene=21&sn=f0a601ca66c01361709e6af8485f629b#wechat_redirect
[CVPR 2021_GAN]: http://mp.weixin.qq.com/s?__biz=MzU5MTgzNzE0MA%3D%3D&chksm=fe2a4289c95dcb9f047c18e63b3be9d78e5e75f1d94dc80be34686ce6b4f22cdee2d7cc7f2ae&idx=1&mid=2247491634&scene=21&sn=427cb6bfd6868662f2b5ae1ac37f6072#wechat_redirect
[CVPR 2021_GAN 1]: http://mp.weixin.qq.com/s?__biz=MzU5MTgzNzE0MA%3D%3D&chksm=fe29bfe4c95e36f275332b31b993532b98d52acc76d26c15e2321de4cf52141510bee3ff96af&idx=1&mid=2247490911&scene=21&sn=671926492c3b36e3350db4963f7fcb16#wechat_redirect
[GAN_StyleGAN]: http://mp.weixin.qq.com/s?__biz=MzU5MTgzNzE0MA%3D%3D&chksm=fe29bedbc95e37cd709f583372918478ad8d3cff5309a85f3a940e7c0356050278f89d7e1c79&idx=2&mid=2247490656&scene=21&sn=bbe156fef077f23263b57e506ad43314#wechat_redirect
[20_ StyleGAN]: http://mp.weixin.qq.com/s?__biz=MzU5MTgzNzE0MA%3D%3D&chksm=fe29b979c95e306f03b134809b16009d06aa0694b03da3368f84625de3700e6a56221d49d1bf&idx=1&mid=2247490498&scene=21&sn=324cfbbde362a728aa5336f4b200f7a3#wechat_redirect
[100_CVPR 2020_GAN]: http://mp.weixin.qq.com/s?__biz=MzU5MTgzNzE0MA%3D%3D&chksm=fe29a94cc95e205ae880cb9518cb16a54c0abde69e5a8654f7fedf70a0cc990775265bf28a1c&idx=1&mid=2247486455&scene=21&sn=94156f53845cdf07e859e343862be745#wechat_redirect
[_ _Python]: http://mp.weixin.qq.com/s?__biz=MzU5MTgzNzE0MA%3D%3D&chksm=fe29abb6c95e22a04d1eb019f964c8ae97b2f3af98c22f730f5d5ed2aea9fdb320b00c688967&idx=2&mid=2247485709&scene=21&sn=55a51c98ab48f393ba305c082ba0e4fe#wechat_redirect
[_ _Think Python]: http://mp.weixin.qq.com/s?__biz=MzU5MTgzNzE0MA%3D%3D&chksm=fe29a418c95e2d0e9a3de46e0e4d1b8d4947154d8639418a511a3091300aa31699932046556e&idx=2&mid=2247485091&scene=21&sn=f569d5d70c2e6c9cc567fbc5147787a4#wechat_redirect
[_ _Pytorch]: http://mp.weixin.qq.com/s?__biz=MzU5MTgzNzE0MA%3D%3D&chksm=fe29aa78c95e236e380c9fd2ae79917d7451ccdd4803c5e6047a16f921d1b61d18efc79ada7c&idx=2&mid=2247485635&scene=21&sn=90608bce32fe353cf3c64b39e3406bc3#wechat_redirect
[_ _2020]: http://mp.weixin.qq.com/s?__biz=MzU5MTgzNzE0MA%3D%3D&chksm=fe29a4a4c95e2db2aab16bdcefa78795bbb7efbc70b533bb197d73f617558f2f8a523120c6e3&idx=2&mid=2247484959&scene=21&sn=563999cc8d48498beed00270844ceec4#wechat_redirect
[Link 1]: http://mp.weixin.qq.com/s?__biz=MzU5MTgzNzE0MA%3D%3D&chksm=fe29a4bac95e2dac758c4aac10a3aac49eb062bc0829d4ef2ebc8e6d5069d246385c5084a2b4&idx=1&mid=2247484929&scene=21&sn=523220858287ed3e6105919703c8e833#wechat_redirect
[_TensorFlow 2.0]: http://mp.weixin.qq.com/s?__biz=MzU5MTgzNzE0MA%3D%3D&chksm=fe29a5b0c95e2ca6599d4a91d5d40900566b7525de3dd85259e7a6c173aa26aa7ea3641d1aaf&idx=2&mid=2247485195&scene=21&sn=e3387067db29d1759755c6c2ea96c870#wechat_redirect
[Link 2]: http://mp.weixin.qq.com/s?__biz=MzU5MTgzNzE0MA%3D%3D&chksm=fe29aa8fc95e23993e057fb0dbe09f2810ea07184a439c173f404a7ab657cadcde3d24efa52a&idx=2&mid=2247485492&scene=21&sn=c4fb135de02cd52f06bc61b0e79476e4#wechat_redirect
[d9ab7c139edeea33a3b5c90e0ebc5b9d.png]: /images/20221005/4ce8d0c62ca64009b9c87a59a703ca52.png