EfficientNet利用训练好的模型测试

EfficientNet 利用训练好的模型测试

Google团队19年的EfficientNet和EfficientDet在图像分类和目标检测方面都取得了惊艳的效果。

1. 训练过程

EfficientNet如何训练自己的数据集我在之前的博文：EfficientNet 训练测试自己的分类数据集介绍过。不再赘述，通过训练我们已经可以在数据集下的model文件夹里保留最好的模型。

2. 测试过程

其实上次的代码稍微修改就可以利用保存的文件进行新的测试，你只需要修改main函数最前面的几行代码，下面放上完整代码：

from __future__ import print_function, division

import torch
import torch.nn as nn
from torch.autograd import Variable
from torchvision import datasets, transforms
import numpy as np
import torch.nn.functional as FUN
import os
from scipy import io

input_size = 224


# Load Test images
def loaddata(data_dir, batch_size, set_name, shuffle):
    data_transforms = {
        'train': transforms.Compose([
            transforms.Resize(input_size),
            transforms.CenterCrop(input_size),
            transforms.RandomAffine(degrees=0, translate=(0.05, 0.05)),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ]),
        'test': transforms.Compose([
            transforms.Resize(input_size),
            transforms.CenterCrop(input_size),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ]),
    }

    image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x), data_transforms[x]) for x in [set_name]}
    # num_workers=0 if CPU else = 1
    dataset_loaders = {x: torch.utils.data.DataLoader(image_datasets[x],
                                                      batch_size=batch_size,
                                                      shuffle=shuffle, num_workers=1) for x in [set_name]}
    data_set_sizes = len(image_datasets[set_name])
    return dataset_loaders, data_set_sizes


def test_model(model, criterion):
    model.eval()
    running_loss = 0.0
    running_corrects = 0
    cont = 0
    outPre = []
    outLabel = []
    dset_loaders, dset_sizes = loaddata(data_dir=data_dir, batch_size=batch_size, set_name='test', shuffle=False)
    for data in dset_loaders['test']:
        inputs, labels = data
        labels = torch.squeeze(labels.type(torch.LongTensor))
        # GPU
        inputs, labels = Variable(inputs.cuda()), Variable(labels.cuda())
        # CPU
        # inputs, labels = Variable(inputs), Variable(labels)
        # forward
        outputs = model(inputs)
        _, preds = torch.max(outputs.data, 1)
        loss = criterion(outputs, labels)
        if cont == 0:
            outPre = outputs.data.cpu()
            outLabel = labels.data.cpu()
        else:
            outPre = torch.cat((outPre, outputs.data.cpu()), 0)
            outLabel = torch.cat((outLabel, labels.data.cpu()), 0)

        running_loss += loss.item() * inputs.size(0)
        running_corrects += torch.sum(preds == labels.data)

        cont += 1

    print('Loss: {:.4f} Acc: {:.4f}'.format(running_loss / dset_sizes,
                                            running_corrects.double() / dset_sizes))

    return FUN.softmax(Variable(outPre)).data.numpy(), outLabel.numpy()


if __name__ == '__main__':
    os.environ["CUDA_VISIBLE_DEVICES"] = "0"
    # Start Testing
    net_name = 'efficientnet-b0'
    data_dir = 'Flowers'
    save_dir = 'Flowers/model'
    modelft_file = save_dir + "/" + net_name + '.pth'
    batch_size = 32

    # GPU时
    model_ft = torch.load(modelft_file).cuda()
    criterion = nn.CrossEntropyLoss().cuda()

    # CPU时
    # model_ft = torch.load(modelft_file, map_location='cpu')
    # criterion = nn.CrossEntropyLoss()

    outPre, outLabel = test_model(model_ft, criterion)

    # Save result
    np.save(save_dir + '/Pre', outPre)
    np.save(save_dir + '/Label', outLabel)

    # Change the result and scores to .mat
    mat = np.load(save_dir + '/Pre.npy')
    io.savemat(save_dir + '/Pre.mat', {'gene_features': mat})

    label = np.load(save_dir + '/Label.npy')
    io.savemat(save_dir + '/Label.mat', {'labels': label})

如果是想测试单张图片也很简单，只需要看一下函数test_model对于单张图片结果的判断，然后稍微修改就好了。