DL学习日记-ViT加载预训练模型并微调

加载预训练模型

方法：可以去torchvision官网上寻找pretrained模型（在左边栏大概中间位置），由于我学习的都是图像方面的，ViT是一个分类模型，我就在右边栏上寻找classification（下图）；在这个包当中有很多经典的模型（可以用作backbone或者baseline都可以，但是作为baseline可能稍微有点老了）

当然我再github上也找到了Google出的ViT模型的一些参数，在ImagNet上训练的，可以参考下面这篇博客ViT pre-trained models 预训练模型下载_vit预训练模型下载-CSDN博客

模型的查看

print(net) # 直接打印整个网络（非递归）

# 递归打印模型模块
for layer in net.modules():
   print(layer)
# 这里可以换成named_modules() 区别就是后者有每个模块的名字（变量名）

#打印方法都是循环打印，不多赘述，下面介绍几个其他函数
net.children() #只打印第一层次的，不显示block的具体细节
net.named_children()

net.state_dict() # 参数名+参数 的一个字典

模型的加载

# 加载预训练模型
model=models.vit_b_16(weights=models.vit_b_16)

#分类头改成9分类问题
num_classes = 9
model.heads.head = nn.Linear(model.heads.head.in_features, num_classes)

# loss function
criterion = nn.CrossEntropyLoss()
# optimizer
optimizer = optim.Adam(model.parameters(), lr=lr, weight_decay=1e-5) #加入权重衰减，防止过拟合
# scheduler
scheduler = StepLR(optimizer, step_size=1, gamma=gamma)

目的：原模型是1000分类问题，我要改成9分类问题

有个问题：我怎么知道我应该改哪个模块

方法：我打印了网络（看过论文之后知道了整个net的结构），然后进行修改。这个方法的前提是充分了解模型的架构了，知道每个block的作用以及参数。

在这里我就是要修改分类头，所以我打印了模型之后找到了最后的heads.head。接下来的loss和优化器学习器不多赘述。由于经过没有预训练模型的效果（训练acc0.8，测试acc0.2），我知道这个模型肯定是过拟合了，所以在源代码的基础上添加了L2正则化（权重衰减）。

模型小插曲-Ubuntu当中的反斜杠与斜杠

我的dataset的网站如下：A Large Scale Fish Dataset | Kaggle

这是一个海洋生物九分类问题，整个数据集是每种海洋生物1000张图片，一共9000张，大小为2.89G。我制作dataset的方法如下代码

import os

import torch
from torch.utils.data import Dataset, DataLoader
from PIL import Image
import os,glob
import csv

from torchvision.transforms import transforms

# 制作csv
datapath='data/FishData/archive'
mydata='data/FishData/mydata.csv'
class_to_num={}
class_name_list=os.listdir(datapath)
# print(class_name_list) ['Black Sea Sprat', 'Gilt-Head Bream', 'Hourse Mackerel', 'Red Mullet',
# 'Sea Bass', 'Shrimp', 'Striped Red Mullet', 'Trout', 'Red Sea Bream']
for class_name in class_name_list:
    class_to_num[class_name]=len(class_to_num.keys())
# print(class_to_num)

image_dir=[]
for class_name in class_name_list:
    image_dir += glob.glob(os.path.join(datapath,class_name,'*.png'))
# print(image_dir)

import random
random.shuffle(image_dir)
with open('data\FishData\mydata.csv',mode='w',newline='') as f:
    writer = csv.writer(f)
    for image in image_dir:
        class_name=image.split(os.sep)[-2]
        label = class_to_num[class_name]
        writer.writerow([image,label])
print('write')

class FishDataset(Dataset):
    def __init__(self, root_dir , resize, mode , transform=None):
        super(FishDataset,self).__init__()
        """
        Args:
            root_dir (string): Directory with all the images.
            transform (callable, optional): Optional transform to be applied on a sample.
        """
        self.root_dir = root_dir
        self.resize=resize
        self.image,self.label=self.load_csv()
        print(f"mode{mode}:{len(self.image)}")
        self.transform=transform
        if self.transform==None:
            self.transform = transforms.Compose([lambda x: Image.open(x).convert('RGB'),
                                                 transforms.Resize(self.resize),
                                                 transforms.ToTensor()])

        if mode=='test':
            self.image=self.image[int(0.8*len(self.image)):]
            self.label = self.label[int(0.8 * len(self.label)):]

        if mode=='val':

            self.image=self.image[int(0.75*len(self.image)):int(0.8*len(self.image))]
            self.label = self.label[int(0.75 * len(self.label)):int(0.8 * len(self.label))]

        if mode=='train':
            self.image = self.image[:int(0.8 * len(self.image))]
            self.label = self.label[:int(0.8 * len(self.label))]

    def load_csv(self):
        image,label=[],[]
        with open(self.root_dir,mode='r') as f:
            reader =csv.reader(f)
            for row in reader:
                i,l=row
                image.append(i)
                label.append(int(l))
        return image ,label

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

    def __getitem__(self, idx):

        img_path = self.image[idx]  # 这应该是图像的路径
        img = Image.open(img_path).convert('RGB')
        image_tensor = self.transform(img)  # 应用 transforms
        label_tensor = torch.tensor(self.label[idx], dtype=torch.long)  # 直接使用 torch.tensor 转换
        return image_tensor, label_tensor

笔者代码能力一般，所以各位海涵

在这里首先是在test，train，val数据集的划分我是train取80%，test取后20%，val取75%-80%这一部分作为验证集。（我在跑代码的时候不小心多打了两个冒号，在切片的代码处，我检查了两天才发现这个错误）

由于我是需要放在服务器上跑代码，所以当我跑的时候报错以下提示：

Traceback (most recent call last):
  File "./train_vit_fish.py", line 144, in <module>
    for data, label in tqdm(train_loader,
  File "/home/shijian/yes/envs/mbyenv/lib/python3.8/site-packages/tqdm/std.py", line 1195, in __iter__
    for obj in iterable:
  File "/home/shijian/yes/envs/mbyenv/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 631, in __next__
    data = self._next_data()
  File "/home/shijian/yes/envs/mbyenv/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 675, in _next_data
    data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
  File "/home/shijian/yes/envs/mbyenv/lib/python3.8/site-packages/torch/utils/data/_utils/fetch.py", line 51, in fetch
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "/home/shijian/yes/envs/mbyenv/lib/python3.8/site-packages/torch/utils/data/_utils/fetch.py", line 51, in <listcomp>
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "/home/shijian/mby/ViT/datapre.py", line 82, in __getitem__
    img = Image.open(img_path)
  File "/home/shijian/yes/envs/mbyenv/lib/python3.8/site-packages/PIL/Image.py", line 3277, in open
    fp = builtins.open(filename, "rb")
FileNotFoundError: [Errno 2] No such file or directory: '/home/shijian/mby/ViT/data/FishData/archive\\Hourse Mackerel\\00246.png'

大概的意思就是找不到文件。我刚看见这个提示第一时间去找了到底有没有这个文件，确实有。然后我在想为什么不是报错00001.png这个文件反而是更大的，所以我当时认为不是系统的问题，而是其他问题，要不然读取第一张图片就会有问题；但是我忽略了我的dataset当中shuffle=true（打乱顺序）的问题，所以就是读取第一张图片就已经出错了。经过查阅我发现是ubuntu系统当中不支持双反斜杠，但是利用glob当中glob方法读取通配符就会自动附加，所以我重写了制作csv文件部分的代码，修改如下

import os
import glob
import csv
import random

datapath='data/FishData/archive'
mydata='data/FishData/mydata.csv'

class_to_num = {}
class_name_list = os.listdir(datapath)
print(class_name_list)

for class_name in class_name_list:
    class_to_num[class_name] = len(class_to_num.keys())

image_dir = []
datapath = f'{datapath}/'
for class_name in class_name_list:
    # 使用os.path.join来确保路径正确
    class_name = f'{class_name}/'
    #print(datapath)
    #print(class_name)
    path=datapath+class_name+'*.png'
    #print(path)
    gpath=os.path.join(path)
    print(gpath)
    images_in_class = glob.glob(gpath)
    print(images_in_class)
    image_dir.extend(images_in_class)

random.shuffle(image_dir)

with open(mydata, mode='w', newline='') as f:
    writer = csv.writer(f)
    for image in image_dir:
        # 使用os.sep获取当前操作系统的路径分隔符
        class_name = os.path.basename(os.path.dirname(image))
        # print(class_name)
        label = class_to_num[class_name]
        writer.writerow([image, label])

print('write')

说实话，细看这个代码，我无非改动了os.path.join部分的代码，其他没变，当我放到服务器上跑的时候路径就正确了，格式都是斜杠而非反斜杠。所以我有个猜测就是这个path的结构跟运行代码的环境有很大关系，但是在这里我也不多探究了。这部分主要是意识到不同系统上不同的读取习惯。