最近用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文件夹