from glob import glob import os import numpy as np import matplotlib.pyplot as plt import shutil from torchvision import transforms from torchvision import models import torch from torch.autograd import Variable import torch.nn as nn from torch.optim import lr_scheduler from torch import optim from torchvision.datasets import ImageFolder from torchvision.utils import make_grid import time def imshow(inp): """Imshow for Tensor.""" 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) path = '../Chapter03/dogsandcats/' files = glob(os.path.join(path,'*/*.jpg'))#读取文件夹内的所有文件，join连接两个或更多的路径名组件 no_of_images = len(files) shuffle = np.random.permutation(no_of_images)#创建可用于创建验证数据集的混合索引 os.mkdir(os.path.join(path,'valid'))#创建保存验证图片集的 validation 目录 os.mkdir(os.path.join(path,'train')) for t in ['train','valid']:#使用标签名称创建目录 for folder in ['dog/','cat/']: os.mkdir(os.path.join(path,t,folder)) for i in shuffle[:2000]:#将图片的一小部分子集复制到 validation 文件夹 #shutil.copy(files[i],'../Chapter03/dogsandcats/valid/') # folder = files[i].split('/')[-1].split('.')[0] # image = files[i].split('/')[-1] # os.rename(files[i],os.path.join(path,'valid',folder,image)) if files[i].split('\\')[-1].split('.')[0]=="cat": shutil.copy(files[i],'../Chapter03/dogsandcats/valid/cat/') else: shutil.copy(files[i],'../Chapter03/dogsandcats/valid/dog/') for i in shuffle[2000:]:#将图片的一小部分子集复制到 training 文件夹 #shutil.copy(files[i],'../chapter3/dogsandcats/train/') # folder = files[i].split('/')[-1].split('.')[0] # image = files[i].split('/')[-1] # os.rename(files[i],os.path.join(path,'train',folder,image)) if files[i].split('\\')[-1].split('.')[0]=="cat": shutil.copy(files[i],'../Chapter03/dogsandcats/train/cat/') else: shutil.copy(files[i],'../Chapter03/dogsandcats/train/dog/')
if torch.cuda.is_available(): is_cuda = True simple_transform = transforms.Compose([transforms.Resize((224,224)) ,transforms.ToTensor() ,transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])]) train = ImageFolder('dogsandcats/train/',simple_transform) valid = ImageFolder('dogsandcats/valid/',simple_transform) train_data_gen = torch.utils.data.DataLoader(train,shuffle=True,batch_size=64,num_workers=3) valid_data_gen = torch.utils.data.DataLoader(valid,batch_size=64,num_workers=3) dataset_sizes = {'train':len(train_data_gen.dataset),'valid':len(valid_data_gen.dataset)} dataloaders = {'train':train_data_gen,'valid':valid_data_gen}
model_ft = models.resnet18(pretrained=True) num_ftrs = model_ft.fc.in_features model_ft.fc = nn.Linear(num_ftrs, 2)#把输出特征改成 2 if torch.cuda.is_available(): model_ft = model_ft.cuda() # Loss and Optimizer learning_rate = 0.001 criterion = nn.CrossEntropyLoss() optimizer_ft = optim.SGD(model_ft.parameters(), lr=0.001, momentum=0.9) exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.1)#StepLR 函数帮助动态修改学习率 def train_model(model, criterion, optimizer, scheduler, num_epochs=5): since = time.time() best_model_wts = model.state_dict() # pytorch 中的 state_dict 是一个简单的python的字典对象,将每一层与它的对应参数建立映射关系.(如model的每一层的weights及偏置等等) # (注意,只有那些参数可以训练的layer才会被保存到模型的state_dict中,如卷积层,线性层等等) # 优化器对象Optimizer也有一个state_dict,它包含了优化器的状态以及被使用的超参数(如lr, momentum,weight_decay等) best_acc = 0.0 for epoch in range(num_epochs): print('Epoch {}/{}'.format(epoch, num_epochs - 1)) print('-' * 10) #每轮都有训练和验证阶段 for phase in ['train', 'valid']: if phase == 'train': scheduler.step() model.train(True) # 模型设为训练模式 else: model.train(False) # 模型设为评估模式 running_loss = 0.0 running_corrects = 0 #在数据上迭代 for data in dataloaders[phase]: # 获取输入 inputs, labels = data # 封装成变量 if torch.cuda.is_available(): inputs = Variable(inputs.cuda()) labels = Variable(labels.cuda()) else: inputs, labels = Variable(inputs), Variable(labels) #梯度参数清 0 optimizer.zero_grad() #前向 outputs = model(inputs) _, preds = torch.max(outputs.data, 1)#torch.max(a,1) 返回每一行中最大值的那个元素，且返回其索引（返回最大元素在这一行的列索引） loss = criterion(outputs, labels) #只在训练阶段反向优化 if phase == 'train': loss.backward() optimizer.step() #统计 running_loss += loss.data running_corrects += torch.sum(preds == labels.data) epoch_loss = running_loss / dataset_sizes[phase] epoch_acc = running_corrects / dataset_sizes[phase] print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc)) #深度复制模型 if phase == 'valid' and epoch_acc > best_acc: best_acc = epoch_acc best_model_wts = model.state_dict() print() time_elapsed = time.time() - since print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60)) print('Best val Acc: {:4f}'.format(best_acc)) #加载最优权重 model.load_state_dict(best_model_wts) return model model_ft = train_model(model_ft, criterion, optimizer_ft, exp_lr_scheduler,num_epochs=5)
以上摘自《PyTorch深度学习》书籍。
