使用VGG16对本地数据集RAF-DB中的basic图片进行训练,官方已经在图片命名时分好了train与test,train和test的label在同一个txt文件里,方便起见,把这两种label分成两个txt;需要自行重写导入数据集的函数。
注意:data.DataLoader后输出的batch中label要变成一维;load训练数据时要打乱顺序
进行50个epoch,最后在测试集中得到83.1%的准确率。
import torch, cv2, os, random
from torch.utils import data
from torchvision import transforms
import torchvision.models as models
import torch.nn as nn
class TxtImage(data.Dataset):
def __init__(self,label,dataRoot,transform=None,size=(120,120),index=0):
self.tranform=transform
self.dataRoot=dataRoot
self.size=size
self.imgList=[]
self.labelList=[]
self.index=index
for xx in label:
x =xx.split(' ')
self.imgList.append(x[0])
self.labelList.append(int(x[1])-1)
def __getitem__(self,index):
imgName=self.imgList[index]
imgName_=list(imgName)
imgName_.insert(self.index,'_aligned')
imgName="".join(imgName_)
imgPath=os.path.join(self.dataRoot,imgName)
img = cv2.imread(imgPath)
if self.size is not None:
img=cv2.resize(img,self.size)
if self.tranform is not None:
img=self.tranform(img)
label=torch.IntTensor([self.labelList[index]])
return img,label
def __len__(self):
return len(self.imgList)
if __name__=='__main__':
transform=transforms.Compose([transforms.ToTensor(),
transforms.Normalize(mean=[0.5,0.5,0.5],std=[0.5,0.5,0.5])])
TestlabelPath = os.path.join('/home/msy/FaceData/RAF-DB/basic/alignByMyself/test.txt')
TrainlabelPath = os.path.join('/home/msy/FaceData/RAF-DB/basic/alignByMyself/train.txt')
ImageRoot='/home/msy/FaceData/RAF-DB/basic/Image/aligned'#aligned
with open(TestlabelPath,'r') as testf:
Testlabels=testf.readlines()
TestDataset=TxtImage(Testlabels,ImageRoot,transform,index=9)
TestLoader=data.DataLoader(TestDataset,batch_size=20,num_workers=2)
TestDataLen=TestDataset.__len__()
with open(TrainlabelPath,'r') as trainf:
Trainlabels=trainf.readlines()
TrainDataset=TxtImage(Trainlabels,ImageRoot,transform,index=11)
TrainLoader=data.DataLoader(TrainDataset,batch_size=20,num_workers=2,shuffle=True)
vgg16 = models.vgg16(pretrained=True)
vgg16.classifier[6] = nn.Linear(4096, 7) #输出改为7个基础表情
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
vgg16.to(device)
loss_function = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(vgg16.parameters(), lr=0.0001)
model_name = "vgg16"
save_path = './{}Net.pth'.format(model_name)
best_acc = 0.0
print("start to train...")
for epoch in range(50):
vgg16.train()
running_loss=0.0
for step,(img,label) in enumerate(TrainLoader):
optimizer.zero_grad()
outputs = vgg16(img.to(device))
label = label.squeeze() # if parameters are none ,tensor is squeezed into one dimension
label = torch.as_tensor(label, dtype=torch.int64, device=device)
loss = loss_function(outputs, label)
loss.backward()
optimizer.step()
# print statistics
running_loss += loss.item()
# print train process
rate = (step + 1) / len(TrainLoader)
a = "*" * int(rate * 50)
b = "." * int((1 - rate) * 50)
print("\rtrain loss: {:^3.0f}%[{}->{}]{:.3f}".format(int(rate * 100), a, b, loss), end="")
print()
# validate
vgg16.eval()
acc = 0.0 # accumulate accurate number / epoch
with torch.no_grad():
for data_test in TestLoader:
test_images, test_labels = data_test
optimizer.zero_grad()
outputs = vgg16(test_images.to(device))
predict_y = torch.max(outputs, dim=1)[1]
test_labels = test_labels.squeeze() # if parameters are none ,tensor is squeezed into one dimension
acc += (predict_y == test_labels.to(device)).sum().item()
accurate_test = round(acc / TestDataLen, 4)
print()
print(acc)
if accurate_test >= best_acc:
best_acc = accurate_test
torch.save(vgg16.state_dict(), save_path)
print("save weights")
print('[epoch %d] train_loss: %.3f test_accuracy: %.3f' %
(epoch + 1, running_loss / step, accurate_test))
print('Finished Training')
本文内容由网友自发贡献,版权归原作者所有,本站不承担相应法律责任。如您发现有涉嫌抄袭侵权的内容,请联系:hwhale#tublm.com(使用前将#替换为@)