[完整案例]编程实现动物头像定位-02
上一篇文章中我们对数据进行了预处理
欲知前情如何，请点击下面的链接
今天我们就来训练我们的模型
首先看一眼我们用到的库
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils import data
import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline
import torchvision
from torchvision import transforms
import os
from lxml import etree
from matplotlib.patches import Rectangle
import glob
from PIL import Image我们之前的图片大小都是不固定的，所以我们需要先resize一下成相同大小的
然后矩形框的位置，我们也改成相对位置
img=pil_img.resize((224,224))
width=int(width)
height=int(height)然后我们重新计算几个坐标的位置，其实就是在对应的轴上的长度，然后再乘以我们的固定长度，也就是以下四行代码的含义，因为需要换成固定大小之后的相对位置
xmin = (xmin/width)*224
ymin = (ymin/height)*224
xmax = (xmax/width)*224
ymax = (ymax/height)*224这个时候我们重新画图来验证一下正确与否
plt.imshow(img)
rect = Rectangle((xmin, ymin), (xmax-xmin), (ymax-ymin), fill=False, color='red')
ax = plt.gca()
ax.axes.add_patch(rect)
plt.show()
没有问题，所以我们继续往下进行
images=glob.glob(r'D:\社交应用\QQ\303\第12章\Oxford-IIIT Pets Dataset\dataset\images\*.jpg')
anno=glob.glob(r'D:\社交应用\QQ\303\第12章\Oxford-IIIT Pets Dataset\dataset\annotations\xmls\*.xml')首先我们读取数据，然后获取对应的标签名
xml_name=[x.split("\\")[-1].replace('.xml','') for x in anno]
imgs=[x for x in images if x.split('\\')[-1].replace('.jpg','') in xml_name]获取与之对应的图片内容
并且定义一个函数获取每张图片的矩形框的相对位置，4个坐标值
def to_labels(path):
xml=open(r'{}'.format(path)).read()
sel=etree.HTML(xml)
width=int(sel.xpath('//size/width/text()')[0])
height=int(sel.xpath('//size/height/text()')[0])
xmin=int(sel.xpath('//bndbox/xmin/text()')[0])
ymin=int(sel.xpath('//bndbox/ymin/text()')[0])
xmax=int(sel.xpath('//bndbox/xmax/text()')[0])
ymax=int(sel.xpath('//bndbox/ymax/text()')[0])
return [xmin/width,ymin/height,xmax/width,ymax/height] #返回比例值labels=[to_labels(p) for p in anno]然后我们手动进行乱序（shuffle）
index=np.random.permutation(len(imgs))
imgs=np.array(imgs)[index]
labels=np.array(labels)[index]这样我们就对我们的数据进行了乱序
然后就是类型转换
labels=labels.astype(np.float32)然后划分训练数据和测试数据
80%的数据作为训练数据
20%的数据作为测试数据
i=int(len(imgs)*0.8)
train_imgs=imgs[:i]
train_labels=labels[:i]
test_imgs=imgs[i:]
test_labels=labels[i:]创建我们的transform
transform=transforms.Compose([
transforms.Resize((224,224)),
transforms.ToTensor(),
])然后创建我们自己的dataset类
class OXford_dataset(data.Dataset):
def __init__(self,img_paths,labels_list):
self.imgs=img_paths
self.labels=labels_list
def __getitem__(self,index):
img=self.imgs[index]
pil_img=Image.open(img)
img_tensor=transform(pil_img)
l1,l2,l3,l4=self.labels[index]
return img_tensor,l1,l2,l3,l4
def __len__(self):
return len(self.imgs)实例化训练dataset和测试dataset
train_dataset=OXford_dataset(train_imgs,train_labels)
test_dataset=OXford_dataset(test_imgs,test_labels)实例化训练dataloader和测试的datalodaer
train_dl=data.DataLoader(train_dataset,batch_size=16,shuffle=True)
test_dl=data.DataLoader(test_dataset,batch_size=16,shuffle=False)需要注意的是测试的dataloader不需要乱序(shuffle）
我们从创建好的dataloader中取2条数据进行绘图看一下是否一切正常
plt.figure(figsize=(12,8))
for i,(img,l1,l2,l3,l4) in enumerate(zip(img_batch[:2],out1_b[:2],out2_b[:2],out3_b[:2],out4_b[:2])):
img=img.permute(1,2,0).numpy()
plt.subplot(1,2,i+1)
plt.imshow(img)
xmin,ymin,xmax,ymax=l1*224,l2*224,l3*224,l4*224
rect=Rectangle((xmin,ymin),(xmax-xmin),(ymax-ymin),fill=False,color='red')
ax=plt.gca()
ax.axes.add_patch(rect)
我们今天训练使用的网络是resnet101，使用预训练好的模型
resnet=torchvision.models.resnet101(pretrained=True)
我们可以打印一下resnet的结构，由于其层数太多，我们就不在这里过多的展示。
in_size=resnet.fc.in_features我们在定义自己的网络的时候只需要使用resnet网络的最后一层即可
class Net(nn.Module):
def __init__(self):
super(Net,self).__init__()
self.conv_base=nn.Sequential(*list(rensnet.children())[:-1])
self.fc1=nn.Linear(in_size,1)
self.fc2=nn.Linear(in_size,1)
self.fc3=nn.Linear(in_size,1)
self.fc4=nn.Linear(in_size,1)
def foward(self,x):
x=self.conv_base(x)
x1=self.fc1(x)
x2=self.fc2(x)
x3=sell.fc3(x)
x4=self.fc4(x)
return x1,x2,x3,x4实例化网络
model = Net()使用gpu，如果有
if torch.cuda.is_available():
model.to('cuda')定义损失函数
loss_fn = nn.MSELoss()根据步长调整学习速率：
from torch.optim import lr_scheduler
optimizer = torch.optim.Adam(model.parameters(), lr=0.0001)
exp_lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)训练函数：
def fit(epoch, model, trainloader, testloader):
total = 0
running_loss = 0
model.train()
for x, y1, y2, y3, y4 in trainloader:
if torch.cuda.is_available():
x, y1, y2, y3, y4 = (x.to('cuda'),
y1.to('cuda'), y2.to('cuda'),
y3.to('cuda'), y4.to('cuda'))
y_pred1, y_pred2, y_pred3, y_pred4 = model(x)
loss1 = loss_fn(y_pred1, y1)
loss2 = loss_fn(y_pred2, y2)
loss3 = loss_fn(y_pred3, y3)
loss4 = loss_fn(y_pred4, y4)
loss = loss1 + loss2 + loss3 + loss4
optimizer.zero_grad()
loss.backward()
optimizer.step()
with torch.no_grad():
running_loss += loss.item()
exp_lr_scheduler.step()
epoch_loss = running_loss / len(trainloader.dataset)
test_total = 0
test_running_loss = 0
model.eval()
with torch.no_grad():
for x, y1, y2, y3, y4 in testloader:
if torch.cuda.is_available():
x, y1, y2, y3, y4 = (x.to('cuda'),
y1.to('cuda'), y2.to('cuda'),
y3.to('cuda'), y4.to('cuda'))
y_pred1, y_pred2, y_pred3, y_pred4 = model(x)
loss1 = loss_fn(y_pred1, y1)
loss2 = loss_fn(y_pred2, y2)
loss3 = loss_fn(y_pred3, y3)
loss4 = loss_fn(y_pred4, y4)
loss = loss1 + loss2 + loss3 + loss4
test_running_loss += loss.item()
epoch_test_loss = test_running_loss / len(testloader.dataset)
print('epoch: ', epoch,
'loss： ', round(epoch_loss, 3),
'test_loss： ', round(epoch_test_loss, 3),
)
return epoch_loss, epoch_test_loss开始训练
epochs = 10
train_loss = []
test_loss = []
for epoch in range(epochs):
epoch_loss, epoch_test_loss = fit(epoch, model, train_dl, test_dl)
train_loss.append(epoch_loss)
test_loss.append(epoch_test_loss)
最后画出我们训练的损失值图像
plt.figure()
plt.plot(range(1, len(train_loss)+1), train_loss, 'r', label='Training loss')
plt.plot(range(1, len(train_loss)+1), test_loss, 'bo', label='Validation loss')
plt.title('Training and Validation Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss Value')
plt.legend()
plt.show()
至此，我们本次的训练就到此结束了。
END
本文参与 腾讯云自媒体同步曝光计划，分享自微信公众号。
原始发表：2022-06-27，如有侵权请联系 cloudcommunity@tencent.com 删除
评论
登录后参与评论
推荐阅读
腾讯云开发者
