一步步完整实现VGGNet分类自己数据集（pytorch ,自己图像数据集准备，前一篇博客的细化）

1 数据准备
1.1 图片路径和标签生成txt文件
1.1.1 原始数据形式如下图

说明：Srinivasan2014文件夹→文件夹下有Train和Test两个子文件夹→每个文件夹下有类别名称为文件夹名的子文件夹→图像数据
1.1.2 生成的路径和形式

1.2 代码 preprocess.py

import os
import glob

import sys 
sys.path.append("..") 
import random

def file_txt():
    BASE ='Srinivasan2014/'
    traindata_path = BASE + '/Train/'
    #类别数
    labels = os.listdir(traindata_path)
    valdata_path =BASE + '/Test/'
    ##写train.txt文件
    txtpath = BASE

    for index, label in enumerate(labels):
        imglist = glob.glob(os.path.join(traindata_path,label, '*.tif'))
        # print(imglist)
        random.shuffle(imglist)
        print(len(imglist))
        trainlist = imglist[:int(0.8*len(imglist))]
        vallist = imglist[(int(0.8*len(imglist))+1):]
        with open(txtpath + 'train.txt', 'a')as f:
            for img in trainlist:
                #第一类赋予标签0，第二类标签1，以此类推
                #保存为（label_file image_label)的形式
                f.write(img + ' ' + str(index))
                f.write('\n')

        with open(txtpath + 'val.txt', 'a')as f:
            for img in vallist:
                f.write(img + ' ' + str(index))
                f.write('\n')

    for index1, label1 in enumerate(labels):
        print(label1)
        imglist1 = glob.glob(os.path.join(valdata_path,label1, '*.tif'))
        random.shuffle(imglist1)
        with open(txtpath + 'test.txt', 'a')as f1:
            for img1 in imglist1:
                f1.write(img1 + ' ' + str(index1))
                f1.write('\n')
    
if __name__ == '__main__':
    file_txt()
    

2 准备喂入神经网络的形式
2.1 代码 data_pre.py

from torch.utils.data import Dataset, DataLoader
import torch
import numpy as np
from PIL import Image
from torchvision import transforms

def transform(size=300):
    train_transform = transforms.Compose([
        transforms.Resize((size, size)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.5,0.5,0.5], std=[0.5,0.5,0.5]),
    ])
    return train_transform

class  SelfCustomDataset(Dataset):
    def __init__(self, label_file):
        
        self.transform=transform()
        with open(label_file, 'r') as f:
            #label_file的格式， （label_file image_label)
            #拆开图像路径和标签
            self.imgs = list(map(lambda line: line.strip().split(' '), f))

    def __getitem__(self, index):
        
        img_path, label = self.imgs[index]
        img = Image.open(img_path).convert('RGB')
        img=self.transform(img)

        return img, torch.from_numpy(np.array(int(label)))
 
    def __len__(self):
        return len(self.imgs)

def load_data():
    batch_size =16
    
    train_label_dir = 'Srinivasan2014/train.txt'
    train_datasets = SelfCustomDataset(train_label_dir)
    train_loader = DataLoader(train_datasets, batch_size=batch_size, shuffle=True, num_workers=0)
    
    val_label_dir ='Srinivasan2014/test.txt'
    val_datasets = SelfCustomDataset(val_label_dir)
    val_loader = DataLoader(val_datasets, batch_size=batch_size, shuffle=True, num_workers=0)
    
    return train_loader,val_loader

if __name__=="__main__":
    train_loader,val_loader=load_data()
    for i, (images, labels) in enumerate(train_loader):
        print('data shape')
        print(images.shape)
        print("***"*10)
        print('label shape')
        print(labels.shape)
        break
    

2.2 输出格式

3代码测试
3.1 VGG网络结构 model.py

from torchvision import models
import torch.nn as nn
class VGGnet(nn.Module):
    def __init__(self,feature_extract=True,num_classes=3):
        super(VGGnet, self).__init__()
        model = models.vgg16(pretrained=True)
        self.features = model.features
        set_parameter_requires_grad(self.features, feature_extract)#固定特征提取层参数
        #自适应输出宽高尺寸为7×7
        self.avgpool=model.avgpool
        self.classifier = nn.Sequential(
            nn.Linear(512*7*7 , 1024),
            nn.ReLU(),
            nn.Linear(1024, 1024),
            nn.ReLU(),
            nn.Linear(1024, num_classes)
        )
        
    def forward(self, x):
        x = self.features(x)
        x = self.avgpool(x)
#        print(x.shape)
        x = x.view(x.size(0), 512*7*7)
        out=self.classifier(x)
        return out
    
def set_parameter_requires_grad(model, feature_extracting):
    if feature_extracting:
        for param in model.parameters():
            param.requires_grad = False
            
if __name__=="__main__":
    import torch
    from torchsummary import summary
    device=torch.device('cuda' if torch.cuda.is_available() else 'cpu' )
    model=VGGnet().to(device)
    summary(model,(3,300,300)) 

3.2 测试代码主函数 main.py

import torch.nn as nn
from model import VGGnet
import torch
from data_pre import load_data

train_loader,test_loader=load_data()

learning_rate=0.001
num_epochs = 2               # train the training data n times, to save time, we just train 1 epoch
LR = 0.01 

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model=VGGnet().to(device)
 
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
# In[]
total_step = len(train_loader)
for epoch in range(num_epochs):
    correct = 0
    total = 0
    for i, (images, labels) in enumerate(train_loader):
#            print(images.shape)
#            break
        images = images.to(device)
        labels = labels.to(device)
 
        # Forward pass
        outputs = model(images)
        loss = criterion(outputs, labels.long())
 
        # Backward and optimize
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
 
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels.long()).sum().item()
        
        if (i + 1) % 2 == 0:
            print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'
                  .format(epoch + 1, num_epochs, i + 1, total_step, loss.item()))
            
        print('Train Accuracy  {} %'.format(100 * correct / total))

model.eval()  #
with torch.no_grad():
    correct = 0
    total = 0
    for images, labels in test_loader:
        images = images.to(device)
        labels = labels.to(device)
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels.long()).sum().item()

    print('Test Accuracy  {} %'.format(100 * correct / total))

3.3测试结果（没用验证集）

4.代码及数据集打包
百度云
提取码：ofmq