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文件夹