trainpy

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 matplotlib.pyplot as plt
from torchvision import datasets,transforms
import torchvision
import numpy as np 
import os
import argparse
import time
from tensorboardX import SummaryWriter

from model.resnet import resnet34
from dataset.DogCat import DogCat

writer = SummaryWriter('resnet34')
since = time.time()
parser=argparse.ArgumentParser()
#argparse是一个Python模块：命令行选项、参数和子命令解析器。
parser.add_argument('--num_workers',type=int,default=0)
#加上--就会输出
parser.add_argument('--batchSize',type=int,default=6)
#一次多少张
parser.add_argument('--nepoch',type=int,default=6)
#训练轮数
parser.add_argument('--lr',type=float,default=0.001)
#学习率
parser.add_argument('--gpu',type=str,default='0')
#-1使用cpu 0123选择gpu
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.RandomAffine(degrees=np.random.randint(0, 180), translate=(0, 0.2), scale=(0.9, 1), shear=(6, 9), fillcolor=66),
	#transforms.RandomHorizontalFlip(),
	#随机水平翻转
	transforms.ToTensor(),
	transforms.Normalize((0.485,0.456,0.406),(0.229,0.224,0.225))
	#归一化
	#input[channel] = (input[channel] - mean[channel]) / std[channel]
	# 三个通道 mean=(0.485,0.456,0.406) std=(0.229,0.224,0.225)
])

#val代表验证集
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_less',transform=transform_train)
valset  =DogCat('./data/train_less',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)
trainset_sizes = {x:len(transform_train[x]) for x in trainloader}
val_sizes = {x:len(transform_train[x]) for x in valloader}
class_names = trainset.label


model=resnet34(pretrained=True)
model.fc=nn.Linear(2048,2)
model.cuda()
optimizer = optim.SGD(model.parameters(), lr=opt.lr, momentum=0.9)
#optimizer=torch.optim.Adam(model.parameters(),lr=opt.lr,betas=(0.9, 0.999), eps=1e-08, weight_decay=0)
#optimizer优化器 随机梯度下降
#优化权重
scheduler=StepLR(optimizer,step_size=3)
#StepLR调整学习率
criterion=nn.CrossEntropyLoss()
#交叉熵
criterion.cuda()


def imshow(inp, title=None):
	# print(inp.size())
	inp = inp.numpy().transpose((1, 2, 0))
	mean = np.array([0.485, 0.456, 0.406])
	std = np.array([0.229, 0.224, 0.225])
	inp = std * inp + mean
	inp = np.clip(inp, 0, 1)
	plt.imshow(inp)
	if title is not None:
		plt.title(title)
	plt.pause(0.001)  # 为了让图像更新可以暂停一会


# Get a batch of training data
inputs, classes = next(iter(trainloader))
# print(inputs.size())
# print(inputs.size())
# Make a grid from batch
out = torchvision.utils.make_grid(inputs)
# print(out.size())
imshow(out,title=[class_names[x] for x in classes])

def train(epoch):
	print('\nEpoch: %d' % epoch)
	running_loss=0.0
	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()))
		running_loss+=loss.item()*batch_idx.size(0)
	epoch_loss=running_loss/trainset_sizes[trainset]
	print("\n")
	print('{}Loss:{:.4f}'.format('train',epoch_loss))
	writer.add_scalar('Train/Loss',epoch_loss,epoch)
	scheduler.step()

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))
	epoch_acc = 1.0 * correct.numpy() / total
	writer.add_scalar('Train/Acc',epoch_acc,epoch)


for epoch in range(opt.nepoch):
	train(epoch)
	val(epoch)

writer.close()
time_elapsed=time.time()-since
print('Training complete in {:.0f}m {:.0f}s'.fomat(time_elapsed//60,time_elapsed%60))


torch.save(model.state_dict(),'ckp/model.pth')