- 🍨 本文为🔗365天深度学习训练营中的学习记录博客
- 🍖 原作者:K同学啊
一.前期准备
1.设置GPU
import torch import torch.nn as nn from torchvision import datasets,transforms import PIL,os,pathlib,warnings device=torch.device('cuda' if torch.cuda.is_available() else 'cpu') device
2.导入数据
data_dir='./data/48-data/' data_dir=pathlib.Path(data_dir) data_path=list(data_dir.glob('*')) classeNames=[str(path).split('\\')[2] for path in data_path] classeNames
train_transforms=transforms.Compose([ transforms.Resize([224,224]), transforms.ToTensor(), transforms.Normalize( mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225] ) ]) total_data=datasets.ImageFolder('./data/48-data/',transform=train_transforms) total_data
total_data.class_to_idx
3.划分数据集
train_size=int(0.8*len(total_data)) test_size=len(total_data)-train_size train_dataset,test_dataset=torch.utils.data.random_split(total_data,[train_size,test_size]) train_dataset,test_dataset
batch_size=32 train_dl=torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=1) test_dl=torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=True, num_workers=1)
for X,y in train_dl: print('shape of X[N,C,H,W]:',X.shape) print('shape of y:',y.shape,y.dtype) break
二.调用官方的VGG-16模型
from torchvision.models import vgg16 model=vgg16(pretrained=True).to(device) for param in model.parameters(): param.requires_grad = False #冻结模型的参数,只训练最后一层的参数 #修改classifier模块的第6层 model.classifier[6]=nn.Linear(4096,len(classeNames)) model.to(device) model
三.训练模型
1.编写训练函数
def train(dataloader,model,loss_fn,optimizer): size=len(dataloader.dataset) num_batches=len(dataloader) train_acc,train_loss=0,0 for X,y in dataloader: X,y=X.to(device),y.to(device) pred=model(X) loss=loss_fn(pred,y) optimizer.zero_grad() loss.backward() optimizer.step() train_acc+=(pred.argmax(1)==y).type(torch.float).sum().item() train_loss+=loss.item() train_acc/=size train_loss/=num_batches return train_acc,train_loss
2.编写测试函数
def test(dataloader,model,loss_fn): size=len(dataloader.dataset) num_batches=len(dataloader) test_loss,test_acc=0,0 with torch.no_grad(): for img,target in dataloader: img,target=img.to(device),target.to(device) target_pred=model(img) loss=loss_fn(target_pred,target) test_loss+=loss.item() test_acc+=(target_pred.argmax(1)==target).type(torch.float).sum().item() test_acc/=size test_loss/=num_batches return test_acc,test_loss
3.设置动态学习率
learn_rate=1e-4 lambda1=lambda epoch: 0.92 ** (epoch // 4) optimizer=torch.optim.Adam(model.parameters(),lr=learn_rate) scheduler=torch.optim.lr_scheduler.LambdaLR(optimizer,lr_lambda=lambda1)
4.正式训练
import copy loss_fn=nn.CrossEntropyLoss() epochs=40 train_acc=[] train_loss=[] test_acc=[] test_loss=[] best_acc=0 for epoch in range(epochs): model.train() epoch_train_acc,epoch_train_loss=train(train_dl,model,loss_fn,optimizer) scheduler.step() model.eval() epoch_test_acc,epoch_test_loss=test(test_dl,model,loss_fn) if epoch_test_acc >best_acc: best_acc=epoch_test_acc best_model=copy.deepcopy(model) train_acc.append(epoch_train_acc) train_loss.append(epoch_train_loss) test_acc.append(epoch_test_acc) test_loss.append(epoch_test_loss) lr=optimizer.state_dict()['param_groups'][0]['lr'] template = ('Epoch:{:2d}, Train_acc:{:.1f}%, Train_loss:{:.3f}, Test_acc:{:.1f}%, Test_loss:{:.3f}, Lr:{:.2E}') print(template.format(epoch+1, epoch_train_acc*100, epoch_train_loss, epoch_test_acc*100, epoch_test_loss, lr)) PATH='./best/best_model.pth' torch.save(model.state_dict(),PATH) print('Done')
四.结果可视化
1.Loss与Accuracy图
import matplotlib.pyplot as plt import warnings warnings.filterwarnings('ignore') plt.rcParams['font.sans-serif']=['SimHei'] plt.rcParams['axes.unicode_minus']=False plt.rcParams['figure.dpi']=100 from datetime import datetime current_time=datetime.now() epochs_range=range(epochs) plt.figure(figsize=(12,3)) plt.subplot(1,2,1) plt.plot(epochs_range,train_acc,label='Training Accuracy') plt.plot(epochs_range,test_acc,label='Test Accuracy') plt.legend(loc='lower right') plt.title('Training and Validation Accuracy') plt.xlabel(current_time) plt.subplot(1,2,2) plt.plot(epochs_range,train_loss,label='Training Loss') plt.plot(epochs_range,test_loss,label='Test Loss') plt.legend(loc='upper right') plt.title('Train and Test Loss') plt.show()
2.指定图片进行预测
from PIL import Image classes=list(total_data.class_to_idx) def predict_one_image(image_path,model,transform,classes): test_img=Image.open(image_path).convert('RGB') plt.imshow(test_img) test_img=transform(test_img) img=test_img.to(device).unsqueeze(0) model.eval() output=model(img) _,pred=torch.max(output,1) pred_class=classes[pred] print(f'预测结果是:{pred_class}')
predict_one_image(image_path='./data/48-data/Angelina Jolie/003_57612506.jpg', model=model, transform=train_transforms, classes=classes)
3.模型评估
best_model.eval() epoch_test_acc,epoch_test_loss=test(test_dl,best_model,loss_fn)
epoch_test_acc,epoch_test_loss
使用了迁移学习的方法,通过加载 VGG16 预训练模型,并在其基础上修改最后一层,使其适应新的分类任务,数据被分为训练集和测试集,并经过归一化等预处理后喂入模型进行训练,优化器采用 Adam,并结合学习率调度器逐步调整学习率,整个训练过程中,模型不断优化,最终选择表现最好的模型进行保存,并通过可视化方式展示训练过程中的准确率和损失变化,同时提供了一个单张图片预测的功能,能够方便地对新图片进行分类预测。