pytorch-搭建图片分类神经网络模板(以Kaggle-ClassifyLeaves为例)

############################### 数据集处理 ###############################
import torch
from torch.utils.data import Dataset,DataLoader
from torchvision import transforms
from torch import nn
from PIL import Image
import os
import matplotlib.pyplot as plt
import numpy as np
import torchvision.models as models
import pandas as pd 
from tqdm import tqdm

train_path = r'E:/Learn_MechineLearning_DeepLearning/pytorch-Learning/Kaggle/Classify-leaves/classify-leaves/train.csv'
label_df = pd.read_csv(train_path)
leave_label = sorted(set(list(label_df['label'])))
n_class = len(leave_label)

class_to_num = dict(zip(leave_label,range(n_class)))
num_to_class = dict(zip(class_to_num.values(),class_to_num.keys()))

############################### Dataset搭建 ###############################
class LeavesData(Dataset):
    def __init__(self , csv_path , img_path , mode = 'train' , valid_ratio = 0.2 , resize_height = 256 , resize_weirgt = 256):
        '''
        Args:
            csv_path(string): csv文件路径 
            img_path(string): 图像文件夹所在路径
            mode(string): 训练模式，测试模式
            valid_ratio(float) : 验证集比例
        '''

        #对原图片尺寸进行统一
        self.resize_height = resize_height
        self.resize_weirgt = resize_weirgt

        self.file_path = img_path
        self.mode = mode

        #读取csv文件
        self.data_info = pd.read_csv(csv_path,header=None) # 不要读取表头
        #计算len
        self.data_len = len(self.data_info.index)-1
        #train_len
        self.train_len = int(self.data_len*(1-valid_ratio))

        if mode == 'train':
            #train图像的名称,label
            self.train_image = np.asarray(self.data_info.iloc[1:self.train_len,0])
            self.train_label = np.asarray(self.data_info.iloc[1:self.train_len,1])
            self.image_arr = self.train_image
            self.labe_arr = self.train_label

        elif mode == 'valid':
            self.valid_image = np.asarray(self.data_info.iloc[self.train_len:,0])
            self.valid_label = np.asarray(self.data_info.iloc[self.train_len:,1])
            self.image_arr = self.valid_image
            self.labe_arr = self.valid_label

        elif mode == 'test':
            self.test_image = np.asarray(self.data_info.iloc[1:,0])
            self.image_arr = self.test_image

        self.real_len = len(self.image_arr)

        print('Finished reading the {} set of Leaves Dataset {} samples found'.format(mode,self.real_len))
        
    def __getitem__(self,index):
        single_image_name  = self.image_arr[index]
        
        # 读到的是文件名，要读取图像文件
        img_as_img = Image.open(self.file_path + single_image_name)
        
        if self.mode == 'train':
            transform = transforms.Compose([
                transforms.Resize((224,224)),
                transforms.RandomHorizontalFlip(p=0.5),
                transforms.RandomVerticalFlip(p=0.5),
                transforms.ToTensor()
            ])
        
        else:
            transform = transforms.Compose([
                transforms.Resize((224,224)),
                transforms.ToTensor()
            ])

        img_as_img = transform(img_as_img)
        
        if self.mode == 'test':
            return img_as_img
        
        else:
            #返回图像label
            label = self.labe_arr[index]
            number_label = class_to_num[label]
        
            return img_as_img , number_label
    
    def __len__(self):
        return self.real_len    
        
train_path = r'E:/Learn_MechineLearning_DeepLearning/pytorch-Learning/Kaggle/Classify-leaves/classify-leaves/train.csv'
test_path = r'E:/Learn_MechineLearning_DeepLearning/pytorch-Learning/Kaggle/Classify-leaves/classify-leaves/test.csv'
img_path = r'E:/Learn_MechineLearning_DeepLearning/pytorch-Learning/Kaggle/Classify-leaves/classify-leaves/'

train_dataset = LeavesData(train_path , img_path , mode = 'train')
val_dataset = LeavesData(train_path , img_path, mode = 'valid')
test_dataset = LeavesData(test_path , img_path , mode = 'test') 

train_loader = torch.utils.data.DataLoader(
    dataset = train_dataset,
    batch_size = 16,
    shuffle = True
)

val_loader = torch.utils.data.DataLoader(
    dataset = val_dataset,
    batch_size = 16,
    shuffle = True
)

test_loader = torch.utils.data.DataLoader(
    dataset = test_dataset,
    batch_size = 16,
    shuffle = True
)

cnt = 0
for img , label in train_loader:
    img = img[0]   # 若出错可改成img = image[0].cuda 试试
    img = img.detach().numpy() # FloatTensor转为ndarray
    img = np.transpose(img, (1,2,0)) # 把channel那一维放到最后
    plt.imshow(img)
    plt.show()
    cnt = cnt+1
    if cnt == 10:
        break


############################### ResNext模型 ###############################
def get_device():
    return 'cuda' if torch.cuda.is_available() else 'cpu'

device = get_device()
print(device)
# 是否要冻住模型的前面一些层
def set_parameter_requires_grad(model, feature_extracting):
#     if feature_extracting:
#         model = model
#         for i, param in enumerate(model.children()):
#             if i == 8:
#                 break
#             param.requires_grad = False
    if feature_extracting:
        model = model
        for param in model.parameters():
            param.requires_grad = False
            
# resnet34模型
def res_model(num_classes, feature_extract = False, use_pretrained=True):

    model_ft = models.resnet34(pretrained=use_pretrained)
    set_parameter_requires_grad(model_ft, feature_extract)
    num_ftrs = model_ft.fc.in_features
#     model_ft.fc = nn.Sequential(
#         nn.Linear(num_ftrs, 512),
#         nn.ReLU(inplace=True),
#         nn.Dropout(.3),
#         nn.Linear(512, len(num_to_class))
#     )
    model_ft.fc = nn.Sequential(
        nn.Linear(num_ftrs, num_classes)
    )
    return model_ft

# resnext50模型
def resnext_model(num_classes, feature_extract = False, use_pretrained=True):

    model_ft = models.resnext50_32x4d(pretrained=use_pretrained)
    set_parameter_requires_grad(model_ft, feature_extract)
    num_ftrs = model_ft.fc.in_features
    model_ft.fc = nn.Sequential(nn.Linear(num_ftrs, num_classes))

    return model_ft

############################### 超参数设置 ###############################
learning_rate = 3e-4
weight_decay = 1e-3
num_epoch = 50
model_path = r'E:/Learn_MechineLearning_DeepLearning/pytorch-Learning/Kaggle/Classify-leaves/classify-leaves/pre_trained_model/pre_trained_resnext50.pthe'



############################### 训练模型 ###############################
# Initialize a model, and put it on the device specified.
#model = res_model(176)
model = resnext_model(176)
model = model.to(device)
model.device = device
# For the classification task, we use cross-entropy as the measurement of performance.
criterion = nn.CrossEntropyLoss()

# Initialize optimizer, you may fine-tune some hyperparameters such as learning rate on your own.
# optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr = learning_rate, weight_decay=weight_decay)
optimizer = torch.optim.Adam(model.parameters(), lr = learning_rate, weight_decay=weight_decay)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=10, eta_min=0, last_epoch=-1)

# The number of training epochs.
n_epochs = num_epoch

best_acc = 0.0
for epoch in range(n_epochs):
    # ---------- Training ----------
    # Make sure the model is in train mode before training.
    model.train() 
    # These are used to record information in training.
    train_loss = []
    train_accs = []
    i = 0
    # Iterate the training set by batches.
    for batch in tqdm(train_loader):
        # A batch consists of image data and corresponding labels.
        imgs, labels = batch
        imgs = imgs.to(device)
        labels = labels.to(device)
        # Forward the data. (Make sure data and model are on the same device.)
        logits = model(imgs)
        # Calculate the cross-entropy loss.
        # We don't need to apply softmax before computing cross-entropy as it is done automatically.
        loss = criterion(logits, labels)
        
        # Gradients stored in the parameters in the previous step should be cleared out first.
        optimizer.zero_grad()
        # Compute the gradients for parameters.
        loss.backward()
        # Update the parameters with computed gradients.
        optimizer.step()
        # 更新学习率------------------------------------------------------------------------------
        scheduler.step()
        if(i % 500 == 0):
            print("learning_rate:", scheduler.get_last_lr()[0])
        i = i + 1
        
        # Compute the accuracy for current batch.
        acc = (logits.argmax(dim=-1) == labels).float().mean()

        # Record the loss and accuracy.
        train_loss.append(loss.item())
        train_accs.append(acc)
        

    
    # The average loss and accuracy of the training set is the average of the recorded values.
    train_loss = sum(train_loss) / len(train_loss)
    train_acc = sum(train_accs) / len(train_accs)

    # Print the information.
    print(f"[ Train | {epoch + 1:03d}/{n_epochs:03d} ] loss = {train_loss:.5f}, acc = {train_acc:.5f}")
    
    # ---------- Validation ----------
    # Make sure the model is in eval mode so that some modules like dropout are disabled and work normally.
    model.eval()
    # These are used to record information in validation.
    valid_loss = []
    valid_accs = []
    
    # Iterate the validation set by batches.
    for batch in tqdm(val_loader):
        imgs, labels = batch
        # We don't need gradient in validation.
        # Using torch.no_grad() accelerates the forward process.
        with torch.no_grad():
            logits = model(imgs.to(device))
            
        # We can still compute the loss (but not the gradient).
        loss = criterion(logits, labels.to(device))

        # Compute the accuracy for current batch.
        acc = (logits.argmax(dim=-1) == labels.to(device)).float().mean()

        # Record the loss and accuracy.
        valid_loss.append(loss.item())
        valid_accs.append(acc)
        
    # The average loss and accuracy for entire validation set is the average of the recorded values.
    valid_loss = sum(valid_loss) / len(valid_loss)
    valid_acc = sum(valid_accs) / len(valid_accs)

    # Print the information.
    print(f"[ Valid | {epoch + 1:03d}/{n_epochs:03d} ] loss = {valid_loss:.5f}, acc = {valid_acc:.5f}")
    
    # if the model improves, save a checkpoint at this epoch
    if valid_acc > best_acc:
        best_acc = valid_acc
        torch.save(model.state_dict(), model_path)
        print('saving model with acc {:.3f}'.format(best_acc))

预训练模型resnext.pth下载