猫狗大战

猫狗大战挑战赛

在 ImageNet 上预训练 的 VGG 网络进行测试。因为原网络的分类结果是1000类，所以这里进行迁移学习，对原网络进行 fine-tune （即固定前面若干层，作为特征提取器，只重新训练最后两层）。

训练过程

训练代码

import numpy as np
import matplotlib.pyplot as plt
import os
import torch
import torch.nn as nn
import torchvision
from torchvision import models,transforms,datasets
import time
import jsondevice = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print('Using gpu: %s ' % torch.cuda.is_available()
!unzip "/content/drive/MyDrive/cat_dog.zip"
#处理数据
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
vgg_format = transforms.Compose([
                 transforms.CenterCrop(224),
                 transforms.ToTensor(),
                 normalize,
               ])
 
data_dir = './cat_dog'
 
dsets = {x: datasets.ImageFolder(os.path.join(data_dir, x), vgg_format)
           for x in ['train', 'val', 'test']}
   
dset_sizes = {x: len(dsets[x]) for x in ['train', 'val', 'test']}
dset_classes = dsets['train'].classes

loader_train = torch.utils.data.DataLoader(dsets['train'], batch_size=128, shuffle=True, num_workers=6)
loader_valid = torch.utils.data.DataLoader(dsets['val'], batch_size=5, shuffle=False, num_workers=6)
loader_test = torch.utils.data.DataLoader(dsets['test'], batch_size=5, shuffle=False, num_workers=6)

!wget https://s3.amazonaws.com/deep-learning-models/image-models/imagenet_class_index.json

model_vgg = models.vgg16(pretrained=True)

print(model_vgg)
 
model_vgg_new = model_vgg;
 
for param in model_vgg_new.parameters():
    param.requires_grad = False

model_vgg_new.classifier._modules['6'] = nn.Linear(4096, 2)
model_vgg_new.classifier._modules['7'] = torch.nn.LogSoftmax(dim = 1)
model_vgg_new = model_vgg_new.to(device)

print(model_vgg_new.classifier)

正式测试

测试代码

import numpy as np
import matplotlib.pyplot as plt
import os
import torch
import torch.nn as nn
import torchvision
from torchvision import models,transforms,datasets
import time
import json


# 判断是否存在GPU设备

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print('Using gpu: %s ' % torch.cuda.is_available())

normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])

vgg_format = transforms.Compose([
                transforms.CenterCrop(224),
                transforms.ToTensor(),
                normalize,
            ])

data_dir = '/content/cat_dog/test'

dsets = {'test': datasets.ImageFolder(data_dir, vgg_format)}
dset_sizes = {x: len(dsets[x]) for x in ['test']}
loader_test = torch.utils.data.DataLoader(dsets['test'], batch_size=1, shuffle=False, num_workers=0)

model_vgg_new = torch.load('/content/drive/MyDrive/models50.9811.pth')
model_vgg_new = model_vgg_new.to(device)


def test(model,dataloader,size):
    model.eval（)
    predictions = np.zeros(size)
    i = 0
    all_preds = {}
    for inputs,classes in dataloader:
        inputs = inputs.to(device)
        outputs = model(inputs)         
        _,preds = torch.max(outputs.data,1)
        # statistics
        #按照batch_size重新分组数据，因为dset中的数据不是按1-2000顺序排列的
        key = dsets['test'].imgs[i][0]
        print(key)
        all_preds[key] = preds[0]
        i += 1
        print('Testing: No. ', i, ' process ... total: ', size)

with open("./drive/MyDrive/result.csv", 'a+') as f:
    for i in range(2000):
        f.write("{},{}\n".format(i, all_preds["./cat_dog/test/TT/"+str(i)+".jpg"]))

test(model_vgg_new,loader_test,size=dset_sizes['test'])