PL-2

PL使用自定义数据训练

导包

from typing import Union, List

import torch
import torch.nn as nn
from PIL import Image
from pytorch_lightning.callbacks import ModelCheckpoint
from pytorch_lightning.core.lightning import LightningModule
from pytorch_lightning.loggers import TensorBoardLogger
from torch.utils.data import DataLoader
from torchsummary import summary
from torchvision import transforms, datasets
import pytorch_lightning as pl
import cv2
import torch.nn.functional as F

定义模型

def Conv2D_BN(in_channels, out_channels, ks, padding, strides=1):
    return nn.Sequential(
        nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=ks, padding=padding, stride=strides),
        nn.BatchNorm2d(out_channels)
    )

def bottleneck_block(in_channels, out_channels=(64, 64, 256), padding=1):
    out_channels_1, out_channels_2, out_channels_3 = out_channels

    base = nn.Sequential()
    base.add_module("bottleneck_block_conv2d_bn1",
                    Conv2D_BN(in_channels, out_channels_1, ks=1, padding=padding, strides=1))
    base.add_module("bottleneck_block_conv2d_bn2", Conv2D_BN(in_channels, out_channels_2, ks=1, padding=1, strides=1))
    base.add_module("bottleneck_block_conv2d_bn3", Conv2D_BN(in_channels, out_channels_3, ks=1, padding=1, strides=1))

    return base

class Net(LightningModule):
    def __init__(self):
        super(Net, self).__init__()
        self.zero_padding = nn.ZeroPad2d(padding=3)
        self.conv_bn = Conv2D_BN(3, 64, ks=7, padding=0, strides=2)
        self.max_pool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)

        self.bottleneck_1 = bottleneck_block(in_channels=64, out_channels=(64, 64, 256), padding=1)
        self.short = Conv2D_BN(in_channels=256, out_channels=256, ks=1, strides=1, padding=0)
        self.avgpool = nn.AvgPool2d(kernel_size=7)
        self.dense = nn.Linear(in_features=256 * 8 * 8, out_features=2)

    def forward(self, x):
        zero_pad = self.zero_padding(x)
        conv_bn = self.conv_bn(zero_pad)
        max_pool = self.max_pool(conv_bn)

        bottleneck_1 = self.bottleneck_1(max_pool)
        short = self.short(bottleneck_1)
        # shortcut = Conv2D_BN(bottleneck_1, 256, ks=1, strides=1, padding=1)
        _x = torch.add(bottleneck_1, short)
        # print(bottleneck_1.shape)
        # print(short.shape)
        _x = self.avgpool(_x)
        _x = _x.view(_x.size(0), -1)
        _x = self.dense(_x)
        return _x

    def training_step(self, batch, batch_idx):
        x, y = batch
        criterion = nn.CrossEntropyLoss()
        loss = criterion(self.forward(x), y)
        self.log('train loss', loss, on_step=True, on_epoch=True, prog_bar=True)
        return loss

    def validation_step(self, batch, batch_idx):
        x, y = batch
        criterion = nn.CrossEntropyLoss()
        loss = criterion(self.forward(x), y)

        return {'val_loss': loss}

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=0.02)

    def train_dataloader(self) -> DataLoader:
        data_transform = transforms.Compose([
            transforms.RandomResizedCrop(224),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                 std=[0.229, 0.224, 0.225])
        ])
        train_data = datasets.ImageFolder(root='dogs-vs-cats/train',
                                          transform=data_transform)
        dataset_loader = torch.utils.data.DataLoader(train_data,
                                                     batch_size=32, shuffle=True,
                                                     num_workers=8)
        return dataset_loader

    def val_dataloader(self) -> Union[DataLoader, List[DataLoader]]:
        data_transform = transforms.Compose([
            transforms.RandomResizedCrop(224),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                 std=[0.229, 0.224, 0.225])
        ])
        train_data = datasets.ImageFolder(root='dogs-vs-cats/test',
                                          transform=data_transform)
        dataset_loader = torch.utils.data.DataLoader(train_data,
                                                     batch_size=32, shuffle=False,
                                                     num_workers=8)
        return dataset_loader

checkpoint_callback = ModelCheckpoint(
        filepath='weights',
        verbose=True,
        monitor='val_loss',
        mode='min'
    )
torch.cuda.empty_cache()
model = Net()
# # print(model)
summary(model, (3, 224, 224))
model.train_dataloader()
trainer = pl.Trainer(max_epochs=10, auto_lr_find=True,
                     auto_scale_batch_size=True, automatic_optimization=True,
                     checkpoint_callback=checkpoint_callback, gpus=1)
trainer.fit(model)

model = model.load_from_checkpoint('weights.ckpt')

GPU available: True, used: True
TPU available: False, using: 0 TPU cores
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]

  | Name         | Type       | Params
--------------------------------------------
0 | zero_padding | ZeroPad2d  | 0     
1 | conv_bn      | Sequential | 9 K   
2 | max_pool     | MaxPool2d  | 0     
3 | bottleneck_1 | Sequential | 25 K  
4 | short        | Sequential | 66 K  
5 | avgpool      | AvgPool2d  | 0     
6 | dense        | Linear     | 32 K  
Epoch 0: val_loss reached 0.87508 (best 0.87508), saving model to D:\MyProjects\torchStu\weights-v0.ckpt as top 1
Epoch 1: val_loss was not in top 1
Epoch 2: val_loss reached 0.85068 (best 0.85068), saving model to D:\MyProjects\torchStu\weights-v1.ckpt as top 1
Epoch 3: val_loss reached 0.11168 (best 0.11168), saving model to D:\MyProjects\torchStu\weights-v0.ckpt as top 1
Epoch 4: val_loss was not in top 1
Epoch 5: val_loss was not in top 1
Epoch 6: val_loss was not in top 1
Epoch 7: val_loss was not in top 1
Epoch 8: val_loss was not in top 1
Epoch 9: val_loss reached 0.00010 (best 0.00010), saving model to D:\MyProjects\torchStu\weights-v1.ckpt as top 1


==========================================================================================
Layer (type:depth-idx)                   Output Shape              Param #
==========================================================================================
├─ZeroPad2d: 1-1                         [-1, 3, 230, 230]         --
├─Sequential: 1-2                        [-1, 64, 112, 112]        --
|    └─Conv2d: 2-1                       [-1, 64, 112, 112]        9,472
|    └─BatchNorm2d: 2-2                  [-1, 64, 112, 112]        128
├─MaxPool2d: 1-3                         [-1, 64, 56, 56]          --
├─Sequential: 1-4                        [-1, 256, 62, 62]         --
|    └─Sequential: 2-3                   [-1, 64, 58, 58]          --
|    |    └─Conv2d: 3-1                  [-1, 64, 58, 58]          4,160
|    |    └─BatchNorm2d: 3-2             [-1, 64, 58, 58]          128
|    └─Sequential: 2-4                   [-1, 64, 60, 60]          --
|    |    └─Conv2d: 3-3                  [-1, 64, 60, 60]          4,160
|    |    └─BatchNorm2d: 3-4             [-1, 64, 60, 60]          128
|    └─Sequential: 2-5                   [-1, 256, 62, 62]         --
|    |    └─Conv2d: 3-5                  [-1, 256, 62, 62]         16,640
|    |    └─BatchNorm2d: 3-6             [-1, 256, 62, 62]         512
├─Sequential: 1-5                        [-1, 256, 62, 62]         --
|    └─Conv2d: 2-6                       [-1, 256, 62, 62]         65,792
|    └─BatchNorm2d: 2-7                  [-1, 256, 62, 62]         512
├─AvgPool2d: 1-6                         [-1, 256, 8, 8]           --
├─Linear: 1-7                            [-1, 2]                   32,770
==========================================================================================
Total params: 134,402
Trainable params: 134,402
Non-trainable params: 0
Total mult-adds (M): 461.60
==========================================================================================
Input size (MB): 0.57
Forward/backward pass size (MB): 49.08
Params size (MB): 0.51
Estimated Total Size (MB): 50.17
==========================================================================================