PyTorch - 迁移学习与实战

迁移学习

迁移学习：迁移学习一般是将大规模数据训练的模型向相关任务上迁移，且这个相关任务只有少量的训练数据。因为大规模数据训练的模型提取的底层特征的泛化能力更高，所以一般冻结底层权重；又因为是不同的任务，所以要放开高层权重。

• 迁移学习一般针对的是目标任务训练数据较少的情况
• 目标训练集的训练样本少，就要注意避免过拟合的问题
• 前几层一般都是来进行特征的提取和抽像

利用迁移学习进行表面缺陷检测

表面缺陷检测数据集

该数据集类的构建与多任务网络实战中无异。

导入包

import torch
import numpy as np
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
import os
import cv2 as cv

对于label编号的对应与解释

# 夹杂 - In - inclusion
# 划痕 - SC - scratch
# 裂纹 - CR- crackle
# 压入氧化皮 - PS - Press in oxide scale
# 麻点 - RS
# 斑点 - PA
defect_labels = ["In","Sc","Cr","PS","RS","Pa"]

_init_()

class SurfaceDefectDataset(Dataset):
    def __init__(self, root_dir):
        self.transform = transforms.Compose([transforms.ToTensor(),
                                             transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                                          std=[0.229, 0.224, 0.225]),
                                             transforms.Resize((200, 200))
                                            ])
        
        img_files = os.listdir(root_dir)
        self.defect_types = []
        self.images = []

        for file_name in img_files:
            defect_attrs = file_name.split("_")
            d_index = defect_labels.index(defect_attrs[0])
            self.images.append(os.path.join(root_dir, file_name