|
|
#coding:utf8
|
|
|
|
|
|
# Copyright 2019 longpeng2008. All Rights Reserved.
|
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
|
# If you find any problem,please contact us
|
|
|
#
|
|
|
# longpeng2008to2012@gmail.com
|
|
|
#
|
|
|
# or create issues
|
|
|
# =============================================================================
|
|
|
from __future__ import print_function, division
|
|
|
|
|
|
import torch
|
|
|
import torch.nn as nn
|
|
|
import torch.optim as optim
|
|
|
from torch.optim import lr_scheduler
|
|
|
import torchvision
|
|
|
from torchvision import datasets, models, transforms
|
|
|
import os
|
|
|
from net import simpleconv3
|
|
|
|
|
|
## 使用tensorboardX进行可视化
|
|
|
from tensorboardX import SummaryWriter
|
|
|
writer = SummaryWriter('logs') ## 创建一个SummaryWriter的示例，默认目录名字为runs
|
|
|
|
|
|
## 训练主函数
|
|
|
def train_model(model, criterion, optimizer, scheduler, num_epochs=25):
|
|
|
for epoch in range(num_epochs):
|
|
|
print('Epoch {}/{}'.format(epoch, num_epochs - 1))
|
|
|
for phase in ['train', 'val']:
|
|
|
if phase == 'train':
|
|
|
model.train(True) ## 设置为训练模式
|
|
|
else:
|
|
|
model.train(False) ## 设置为验证模式
|
|
|
|
|
|
running_loss = 0.0 ##损失变量
|
|
|
running_accs = 0.0 ##精度变量
|
|
|
number_batch = 0 ##
|
|
|
|
|
|
## 从dataloaders中获得数据
|
|
|
for data in dataloaders[phase]:
|
|
|
inputs, labels = data
|
|
|
if use_gpu:
|
|
|
inputs = inputs.cuda()
|
|
|
labels = labels.cuda()
|
|
|
|
|
|
optimizer.zero_grad() ##清空梯度
|
|
|
outputs = model(inputs) ##前向运行
|
|
|
_, preds = torch.max(outputs.data, 1) ##使用max()函数对输出值进行操作，得到预测值索引
|
|
|
loss = criterion(outputs, labels) ##计算损失
|
|
|
if phase == 'train':
|
|
|
loss.backward() ##误差反向传播
|
|
|
optimizer.step() ##参数更新
|
|
|
scheduler.step() ## 调整学习率
|
|
|
|
|
|
running_loss += loss.data.item()
|
|
|
running_accs += torch.sum(preds == labels).item()
|
|
|
number_batch += 1
|
|
|
|
|
|
## 得到每一个epoch的平均损失与精度
|
|
|
epoch_loss = running_loss / number_batch
|
|
|
epoch_acc = running_accs / dataset_sizes[phase]
|
|
|
|
|
|
## 收集精度和损失用于可视化
|
|
|
if phase == 'train':
|
|
|
writer.add_scalar('data/trainloss', epoch_loss, epoch)
|
|
|
writer.add_scalar('data/trainacc', epoch_acc, epoch)
|
|
|
else:
|
|
|
writer.add_scalar('data/valloss', epoch_loss, epoch)
|
|
|
writer.add_scalar('data/valacc', epoch_acc, epoch)
|
|
|
|
|
|
print('{} Loss: {:.4f} Acc: {:.4f}'.format(
|
|
|
phase, epoch_loss, epoch_acc))
|
|
|
|
|
|
writer.close()
|
|
|
return model
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
|
|
|
crop_size = 48 ##图像裁剪大小，即训练输入大小
|
|
|
nclass = 4 ##分类类别数
|
|
|
model = simpleconv3(nclass) ##创建模型
|
|
|
data_dir = './data' ##数据目录
|
|
|
|
|
|
## 模型缓存接口
|
|
|
if not os.path.exists('models'):
|
|
|
os.mkdir('models')
|
|
|
|
|
|
## 检查GPU是否可用，如果是使用GPU，否使用CPU
|
|
|
use_gpu = torch.cuda.is_available()
|
|
|
if use_gpu:
|
|
|
model = model.cuda()
|
|
|
print(model)
|
|
|
|
|
|
## 创建数据预处理函数，训练预处理包括随机裁剪缩放、随机翻转、归一化，验证预处理包括中心裁剪，归一化
|
|
|
data_transforms = {
|
|
|
'train': transforms.Compose([
|
|
|
transforms.RandomResizedCrop(crop_size),
|
|
|
transforms.RandomHorizontalFlip(),
|
|
|
transforms.RandomRotation(20),
|
|
|
transforms.ToTensor(),
|
|
|
transforms.Normalize([0.5,0.5,0.5], [0.5,0.5,0.5])
|
|
|
]),
|
|
|
'val': transforms.Compose([
|
|
|
transforms.Resize(crop_size),
|
|
|
transforms.CenterCrop(crop_size),
|
|
|
transforms.ToTensor(),
|
|
|
transforms.Normalize([0.5,0.5,0.5], [0.5,0.5,0.5])
|
|
|
]),
|
|
|
}
|
|
|
|
|
|
## 使用torchvision的dataset ImageFolder接口读取数据
|
|
|
image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x),
|
|
|
data_transforms[x]) for x in ['train', 'val']}
|
|
|
|
|
|
## 创建数据指针，设置batch大小，shuffle，多进程数量
|
|
|
dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x],
|
|
|
batch_size=96,
|
|
|
shuffle=True,
|
|
|
num_workers=4) for x in ['train', 'val']}
|
|
|
## 获得数据集大小
|
|
|
dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']}
|
|
|
|
|
|
## 优化目标使用交叉熵，优化方法使用带动量项的SGD，学习率迭代策略为step，每隔100个epoch，变为原来的0.1倍
|
|
|
criterion = nn.CrossEntropyLoss()
|
|
|
optimizer_ft = optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
|
|
|
step_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=100, gamma=0.1)
|
|
|
|
|
|
model = train_model(model=model,
|
|
|
criterion=criterion,
|
|
|
optimizer=optimizer_ft,
|
|
|
scheduler=step_lr_scheduler,
|
|
|
num_epochs=300)
|
|
|
|
|
|
torch.save(model.state_dict(),'models/model.pt') |
