365天深度学习训练营-第P6周：好莱坞明星识别

• 🍨 本文为🔗365天深度学习训练营 中的学习记录博客
• 🍖 原作者：K同学啊

 我的环境：

• 语言环境：Python3.11.2
• 编译器：PyCharm Community Edition 2022.3
• 深度学习环境：torch==2.0,torchvision==0.15.1 

 本次目标：

• 调用官方vgg16模型

---

 一、前期准备

        将数据集导入。

import torch
import torch.nn as nn
from torchvision import transforms,datasets
from torchvision.models import vgg16
import PIL,pathlib
from torchinfo import summary
import matplotlib.pyplot as plt

data_path='F:\\48-data'
data_path=pathlib.Path(data_path)
data_paths = list(data_path.glob('*'))
classNames = [str(path).split('\\')[2] for path in data_paths]
print(classNames)

'''
['Angelina Jolie', 'Brad Pitt', 'Denzel Washington', 'Hugh Jackman', 'Jennifer Lawrence', 'Johnny Depp', 'Kate Winslet', 'Leonardo DiCaprio', 'Megan Fox', 'Natalie Portman', 'Nicole Kidman', 'Robert Downey Jr', 'Sandra Bullock', 'Scarlett Johansson', 'Tom Cruise', 'Tom Hanks', 'Will Smith']
'''

        处理数据并划分数据集。

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_path, transform=train_transforms)
print(total_data)
print(total_data.class_to_idx)

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])
print(train_dataset,test_dataset)
print(train_size,test_size)

batch_size = 32
train_dl = torch.utils.data.DataLoader(train_dataset,batch_size=batch_size,shuffle=True)
test_dl = torch.utils.data.DataLoader(test_dataset,batch_size=batch_size)

二、调用官方vgg16模型

model = vgg16(pretrained=True).to("cpu")
for params in model.parameters():
    params.requires_grad=False
model.classifier._modules['6']=nn.Linear(4096,len(classNames))
summary(model)

        调用官方vgg16模型进行学习。

​

三、训练模型

        设置超参数。

loss_fn = nn.CrossEntropyLoss()
learn_rate = 1e-3
opt = torch.optim.SGD(model.parameters(),lr=learn_rate)

        编写训练函数和测试函数。

def train(dataloader,model,loss_fn,optimizer):
    size = len(dataloader.dataset)
    num_batchs = len(dataloader)
    train_loss, train_acc = 0, 0
    for X,y in dataloader:
        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_batchs

    return train_acc,train_loss

def test(dataloader,model,loss_fn):
    size=len(dataloader.dataset)
    num_batchs=len(dataloader)
    test_loss,test_acc=0,0

    with torch.no_grad():
        for imgs,target in dataloader:

            target_pred = model(imgs)
            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_batchs

    return test_acc,test_loss

        正式训练。

epochs=20
train_loss=[]
train_acc=[]
test_loss=[]
test_acc=[]

for epoch in range(epochs):
    model.train()
    epoch_train_acc,epoch_train_loss= train(train_dl,model,loss_fn,opt)
    model.eval()
    epoch_test_acc,epoch_test_loss=test(test_dl,model, loss_fn)
    train_acc.append(epoch_train_acc)
    train_loss.append(epoch_train_loss)
    test_acc.append(epoch_test_acc)
    test_loss.append(epoch_test_loss)

    template = ('Epoch:{:2d},Train_acc:{:.1f}%,Train_loss:{:.3f},Test_acc:{:.1f}%,Test_loss:{:.3f}')
    print(template.format(epoch+1,epoch_train_acc,epoch_train_loss,epoch_test_acc,epoch_test_loss))
print('Done')

​

四、结果可视化

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.subplot(1,2,2)
plt.plot(epochs_range,train_loss,label='Training Loss')
plt.plot(epochs_range,test_loss,label='Validation Loss')
plt.legend(loc='upper right')
plt.title('Training and Validation Loss')
plt.show()

​

五、指定图片进行预测 

classes = list(total_data.class_to_idx)
def predict_one_img(image_path, model, transform, classes):
    test_img = PIL.Image.open(image_path).convert('RGB')
    test_img=transform(test_img)
    img = test_img.to('cpu').unsqueeze(0)
    model.eval()
    output = model(img)
    x,pred = torch.max(output,1)
    pred_class=classes[pred]
    print(f'预测结果是{pred_class}')
predict_one_img(image_path='F:\\48-data\\Angelina Jolie\\001_fe3347c0.jpg',model=model,transform=train_transforms,classes=classes)

        保存模型参数

PATH = './model.pth'
torch.save(model.state_dict(),PATH)

model.load_state_dict(torch.load(PATH))