使用VGG模型进行猫狗大战

下载数据

使用课程提供的数据下载地址：http://fenggao-image.stor.sinaapp.com/dogscats.zip
解压数据

数据预处理

打印预览数据

创建 VGG Model

使用预训练好的 VGG 模型。将最后一层进行改写nn.Linear(4096, 2)使得最后输出的结果只有两个，即分辨猫与狗。

训练并测试全连接层

使用Adam优化器对模型进行优化

'''
第一步：创建损失函数和优化器

损失函数 NLLLoss() 的 输入 是一个对数概率向量和一个目标标签. 
它不会为我们计算对数概率，适合最后一层是log_softmax()的网络. 
'''
criterion = nn.NLLLoss()

# 学习率
lr = 0.001
# 这里使用Adam优化器
optimizer_vgg = torch.optim.Adam(model_vgg_new.classifier[6].parameters(),lr = lr)

'''
第二步：训练模型并保存
model: 训练的模型
dataloader: 训练集
size: 训练集大小
epochs: 训练次数
optimizer: 优化器
'''
def train_model(model,dataloader,size,epochs=1,optimizer=None):
  model.train()  #用于模型训练
  
  for epoch in range(epochs):
    epoch_acc_max = 0
    running_loss = 0.0
    running_corrects = 0
    count = 0

    for inputs,classes in dataloader:
      inputs = inputs.to(device)
      classes = classes.to(device)
            
      outputs = model(inputs)	#参数前向传播
            
      loss = criterion(outputs,classes)           
      optimizer = optimizer
      optimizer.zero_grad()	#优化器梯度初始化
      loss.backward()			#梯度反向传播
      optimizer.step()
      _,preds = torch.max(outputs.data,1) #得到预测结果
        # statistics
      running_loss += loss.data.item()
      running_corrects += torch.sum(preds == classes.data)
                    
      count += len(inputs)
      print('Training: No. ', count, ' process ... total: ', size)
            
    epoch_loss = running_loss / size
    epoch_acc = running_corrects.data.item() / size
        
    if epoch_acc > epoch_acc_max:
	    epoch_acc_max = epoch_acc
	    torch.save(model, 'model_best.pth')	#保存最好模型
        
    print('Loss: {:.4f} Acc: {:.4f}'.format(
                     epoch_loss, epoch_acc))
        
        
# 模型训练
train_model(model_vgg_new, loader_train,size = dset_sizes['train'], 
		epochs = 5, optimizer=optimizer_vgg)

测试test