VGG16训练RAF-DB

使用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')