kaggle比赛: Dogs_vs_Cats之PyTorch实现

最近用pytorch写了一下kaggle入门的比赛：Dogs_vs_Cats

代码在个人github:

https://github.com/JackwithWilshere/Kaggle-Dogs_vs_Cats_PyTorch

1.数据data

数据处理部分用的pytorch的Dataset,

DogCat.py代码如下：

import os
import random
from PIL import Image
import torch.utils.data as data
import numpy as np 
import torchvision.transforms as transforms

class DogCat(data.Dataset):
	def __init__(self,root,transform=None,train=True,test=False):
		self.test=test
		self.train=train
		self.transform=transform
		imgs=[os.path.join(root,img)for img in os.listdir(root)]

		# test1: data/test1/8973.jpg
		# train: data/train/cat.10004.jpg

		if self.test:
			imgs=sorted(imgs,key=lambda x: int(x.split('.')[-2].split('/')[-1]))
		else:
			imgs=sorted(imgs,key=lambda x: int(x.split('.')[-2]))

		imgs_num=len(imgs)
		
		if self.test:
			self.imgs=imgs
		else:
			random.shuffle(imgs)
			if self.train:
				self.imgs=imgs[:int(0.7*imgs_num)]
			else:
				self.imgs=imgs[int(0.7*imgs_num):]
				
	def __getitem__(self,index):
		img_path=self.imgs[index]
		if self.test:
			label=int(self.imgs[index].split('.')[-2].split('/')[-1])
		else:
			label=1 if 'dog' in img_path.split('/')[-1] else 0
		data=Image.open(img_path)
		data=self.transform(data)
		return data,label
	def __len__(self):
return len(self.imgs)

2.模型model

直接选用的resnet

3.训练train

train.py代码如下

import torch
import torch.nn as nn
import torch.optim as optim
from torch.autograd import Variable
from torch.optim.lr_scheduler import *
import torchvision.transforms as transforms
import numpy as np 
import os
import argparse

from model.resnet import resnet101
from dataset.DogCat import DogCat

parser=argparse.ArgumentParser()
parser.add_argument('--num_workers',type=int,default=2)
parser.add_argument('--batchSize',type=int,default=64)
parser.add_argument('--nepoch',type=int,default=21)
parser.add_argument('--lr',type=float,default=0.001)
parser.add_argument('--gpu',type=str,default='7')
opt=parser.parse_args()
print(opt)
os.environ["CUDA_VISIBLE_DEVICES"]=opt.gpu

transform_train=transforms.Compose([
	transforms.Resize((256,256)),
	transforms.RandomCrop((224,224)),
	transforms.RandomHorizontalFlip(),
	transforms.ToTensor(),
	transforms.Normalize((0.485,0.456,0.406),(0.229,0.224,0.2225))
])

transform_val=transforms.Compose([ 
	transforms.Resize((224,224)),
	transforms.ToTensor(),
	transforms.Normalize((0.485,0.456,0.406),(0.229,0.224,0.225)),
])

trainset=DogCat('./data/train',transform=transform_train)
valset  =DogCat('./data/train',transform=transform_val)
trainloader=torch.utils.data.DataLoader(trainset,batch_size=opt.batchSize,shuffle=True,num_workers=opt.num_workers)
valloader=torch.utils.data.DataLoader(valset,batch_size=opt.batchSize,shuffle=False,num_workers=opt.num_workers)

model=resnet101(pretrained=True)
model.fc=nn.Linear(2048,2)
model.cuda()
optimizer=torch.optim.SGD(model.parameters(),lr=opt.lr,momentum=0.9,weight_decay=5e-4)
scheduler=StepLR(optimizer,step_size=3)
criterion=nn.CrossEntropyLoss()
criterion.cuda()

def train(epoch):
	print('\nEpoch: %d' % epoch)
	scheduler.step()
	model.train()
	for batch_idx,(img,label) in enumerate(trainloader):
		image=Variable(img.cuda())
		label=Variable(label.cuda())
		optimizer.zero_grad()
		out=model(image)
		loss=criterion(out,label)
		loss.backward()
		optimizer.step()
		print("Epoch:%d [%d|%d] loss:%f" %(epoch,batch_idx,len(trainloader),loss.mean()))

def val(epoch):
	print("\nValidation Epoch: %d" %epoch)
	model.eval()
	total=0
	correct=0
	with torch.no_grad():
		for batch_idx,(img,label) in enumerate(valloader):
			image=Variable(img.cuda())
			label=Variable(label.cuda())
			out=model(image)
			_,predicted=torch.max(out.data,1)
			total+=image.size(0)
			correct+=predicted.data.eq(label.data).cpu().sum()
	print("Acc: %f "% ((1.0*correct.numpy())/total))

for epoch in range(opt.nepoch):
	train(epoch)
	val(epoch)
torch.save(model.state_dict(),'ckp/model.pth')

4.test.py

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.transforms as transforms
import os
from torch.autograd import Variable
import argparse
import numpy as np 
from torch.optim.lr_scheduler import *
import csv

from model.resnet import resnet101
from dataset.DogCat import DogCat 

parser = argparse.ArgumentParser()
parser.add_argument('--num_workers', type=int, default=2)
parser.add_argument('--batchSize', type=int, default=64)
parser.add_argument('--nepoch', type=int, default=1, help='number of epochs to train for')
parser.add_argument('--lr', type=float, default=0.0002, help='learning rate, default=0.0002')
parser.add_argument('--gpu', type=str, default='7', help='gpu ids: e.g. 0  0,1,2, 0,2. use -1 for CPU')
opt = parser.parse_args()
print(opt)
os.environ["CUDA_VISIBLE_DEVICES"] = opt.gpu

transform_test = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
])

testset=DogCat('./data/test1',transform=transform_test,train=False,test=True)
testloader=torch.utils.data.DataLoader(testset,batch_size=opt.batchSize,shuffle=False,num_workers=opt.num_workers)
model=resnet101(pretrained=True)
model.fc=nn.Linear(2048,2)
model.load_state_dict(torch.load('ckp/model.pth'))
model.cuda()
model.eval()
results=[]

with torch.no_grad():
    for image,label in testloader:
        image=Variable(image.cuda())
        out=model(image)
        label=label.numpy().tolist()
        _,predicted=torch.max(out.data,1)
        predicted=predicted.data.cpu().numpy().tolist()
        results.extend([[i,";".join(str(j))] for (i,j) in zip(label,predicted)])


eval_csv=os.path.join(os.path.expanduser('.'),'deploy','eval.csv')

with open(eval_csv,'w',newline='') as f:
    writer=csv.writer(f,delimiter=',')
    q=("id","label")
    writer.writerow(q)
    for x in results:
writer.writerow(x)

最后生成的eval文件在deploy文件夹