ImageDataset.py
# -*- coding: utf-8 -*-
# @Time : 2021/1/28 17:23
# @Author : Johnson
'''
负责训练以及测试数据的读取
'''
from torchvision import transforms,datasets
import os
import torch
from PIL import Image
def readImg(path):
'''
用于替代ImageFolder的默认读取图片函数，以读取单通道图片
:param path:
:return:
'''
return Image.open(path)
def ImageDataset(args):
#数据增强以及归一化
#图片都是190*100，训练时随机裁剪90*90，测试时裁剪中间的90*90
data_transforms = {
'train': transforms.Compose([
transforms.RandomCrop(90),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize([0.5], [0.5])
]),
'test': transforms.Compose([
transforms.CenterCrop(90),
transforms.ToTensor(),
transforms.Normalize([0.5], [0.5])
]),
}
data_dir = args.data_dir
image_datasets = {x:datasets.ImageFolder(os.path.join(data_dir,x),data_transforms[x],loader=readImg) for x in ['train','test']}
dataloaders = {x:torch.utils.data.DataLoader(image_datasets[x],batch_size=args.batch_size,shuffle=(x=='train'),num_workers=args.num_workers) for x in ['train','test']}
dataset_sizes = {x:len(image_datasets[x]) for x in ['train','test']}
class_names = image_datasets['train'].classes
return dataloaders,dataset_sizes,class_names
simpleNet.py
# -*- coding: utf-8 -*-
# @Time : 2021/1/28 17:35
# @Author : Johnson
'''
简单的用于分类的网络
'''
import torch
import torch.nn as nn
class SimpleNet(nn.Module):
def __init__(self):
super(SimpleNet, self).__init__()
# 三个卷积层用于提取特征
# 1 input channel image 90x90, 8 output channel image 44x44
self.conv1 = nn.Sequential(
nn.Conv2d(in_channels=1, out_channels=8, kernel_size=3, stride=1, padding=0),
nn.ReLU(),
nn.MaxPool2d(2)
)
# 8 input channel image 44x44, 16 output channel image 22x22
self.conv2 = nn.Sequential(
nn.Conv2d(in_channels=8, out_channels=16, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.MaxPool2d(2)
)
# 16 input channel image 22x22, 32 output channel image 10x10
self.conv3 = nn.Sequential(
nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, stride=1, padding=0),
nn.ReLU(),
nn.MaxPool2d(2)
)
# 分类
self.classifier = nn.Sequential(
nn.Linear(32 * 10 * 10, 3)
)
def forward(self, x):
x = self.conv1(x)
x = self.conv2(x)
x = self.conv3(x)
x = x.view(-1, 32 * 10 * 10)
x = self.classifier(x)
return x
train.py
# -*- coding: utf-8 -*-
# @Time : 2021/1/28 17:35
# @Author : Johnson
'''
进行训练
'''
import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
import time
import os
import json
from math import ceil
import argparse
import copy
from ImageDataset import ImageDataset
from simpleNet import SimpleNet
from tensorboardX import SummaryWriter
writer = SummaryWriter(log_dir='log')
def train_model(args, model, criterion, optimizer, scheduler, num_epochs, dataset_sizes, use_gpu):
since = time.time()
best_model_wts = copy.deepcopy(model.state_dict())
best_acc = 0.0
device = torch.device('cuda' if use_gpu else 'cpu')
for epoch in range(args.start_epoch, num_epochs):
# 每一个epoch中都有一个训练和一个验证过程(Each epoch has a training and validation phase)
for phase in ['train', 'test']:
if phase == 'train':
scheduler.step(epoch)
# 设置为训练模式(Set model to training mode)
model.train()
else:
# 设置为验证模式(Set model to evaluate mode)
model.eval()
running_loss = 0.0
running_corrects = 0
tic_batch = time.time()
# 在多个batch上依次处理数据(Iterate over data)
for i, (inputs, labels) in enumerate(dataloders[phase]):
inputs = inputs.to(device)
labels = labels.to(device)
# 梯度置零(zero the parameter gradients)
optimizer.zero_grad()
# 前向传播(forward)
# 训练模式下才记录梯度以进行反向传播(track history if only in train)
with torch.set_grad_enabled(phase == 'train'):
outputs = model(inputs)
_, preds = torch.max(outputs, 1)
loss = criterion(outputs, labels)
# 训练模式下进行反向传播与梯度下降(backward + optimize only if in training phase)
if phase == 'train':
loss.backward()
optimizer.step()
# 统计损失和准确率(statistics)
running_loss += loss.item() * inputs.size(0)
running_corrects += torch.sum(preds == labels.data)
batch_loss = running_loss / (i * args.batch_size + inputs.size(0))
batch_acc = running_corrects.double() / (i * args.batch_size + inputs.size(0))
if phase == 'train' and (i + 1) % args.print_freq == 0:
print(
'[Epoch {}/{}]-[batch:{}/{}] lr:{:.6f} {} Loss: {:.6f} Acc: {:.4f} Time: {:.4f} sec/batch'.format(
epoch + 1, num_epochs, i + 1, ceil(dataset_sizes[phase] / args.batch_size),
scheduler.get_lr()[0], phase, batch_loss, batch_acc,
(time.time() - tic_batch) / args.print_freq))
tic_batch = time.time()
epoch_loss = running_loss / dataset_sizes[phase]
epoch_acc = running_corrects.double() / dataset_sizes[phase]
if epoch == 0:
os.remove('result.txt')
with open('result.txt', 'a') as f:
f.write('Epoch:{}/{} {} Loss: {:.4f} Acc: {:.4f} \n'.format(epoch + 1, num_epochs, phase, epoch_loss,
epoch_acc))
print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))
writer.add_scalar(phase + '/Loss', epoch_loss, epoch)
writer.add_scalar(phase + '/Acc', epoch_acc, epoch)
if (epoch + 1) % args.save_epoch_freq == 0:
if not os.path.exists(args.save_path):
os.makedirs(args.save_path)
torch.save(model.state_dict(), os.path.join(args.save_path, "epoch_" + str(epoch) + ".pth"))
# 深拷贝模型(deep copy the model)
if phase == 'test' and epoch_acc > best_acc:
best_acc = epoch_acc
best_model_wts = copy.deepcopy(model.state_dict())
# 将model保存为graph
writer.add_graph(model, (inputs,))
time_elapsed = time.time() - since
print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))
print('Best val Accuracy: {:4f}'.format(best_acc))
# 载入最佳模型参数(load best model weights)
model.load_state_dict(best_model_wts)
return model
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='classification')
# 图片数据的根目录(Root catalog of images)
parser.add_argument('--data-dir', type=str, default='images')
parser.add_argument('--batch-size', type=int, default=16)
parser.add_argument('--num-epochs', type=int, default=150)
parser.add_argument('--lr', type=float, default=0.045)
parser.add_argument('--num-workers', type=int, default=4)
parser.add_argument('--print-freq', type=int, default=1)
parser.add_argument('--save-epoch-freq', type=int, default=1)
parser.add_argument('--save-path', type=str, default='output')
parser.add_argument('--resume', type=str, default='', help='For training from one checkpoint')
parser.add_argument('--start-epoch', type=int, default=0, help='Corresponding to the epoch of resume')
args = parser.parse_args()
# read data
dataloders, dataset_sizes, class_names = ImageDataset(args)
with open('class_names.json', 'w') as f:
json.dump(class_names, f)
# use gpu or not
use_gpu = torch.cuda.is_available()
print("use_gpu:{}".format(use_gpu))
# get model
model = SimpleNet()
if args.resume:
if os.path.isfile(args.resume):
print(("=> loading checkpoint '{}'".format(args.resume)))
model.load_state_dict(torch.load(args.resume))
else:
print(("=> no checkpoint found at '{}'".format(args.resume)))
if use_gpu:
model = torch.nn.DataParallel(model)
model.to(torch.device('cuda'))
else:
model.to(torch.device('cpu'))
# 用交叉熵损失函数(define loss function)
criterion = nn.CrossEntropyLoss()
# 梯度下降(Observe that all parameters are being optimized)
optimizer_ft = optim.SGD(model.parameters(), lr=args.lr, momentum=0.9, weight_decay=0.00004)
# Decay LR by a factor of 0.98 every 1 epoch
exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=1, gamma=0.98)
model = train_model(args=args,
model=model,
criterion=criterion,
optimizer=optimizer_ft,
scheduler=exp_lr_scheduler,
num_epochs=args.num_epochs,
dataset_sizes=dataset_sizes,
use_gpu=use_gpu)
torch.save(model.state_dict(), os.path.join(args.save_path, 'best_model_wts.pth'))
writer.close()
test.py
# -*- coding: utf-8 -*-
# @Time : 2021/1/28 17:36
# @Author : Johnson
'''
测试分类
'''
from PIL import Image
from torchvision import transforms
import torch
from torch.autograd import Variable
import os
import json
from simpleNet import SimpleNet
def predict_image(model, image_path):
image = Image.open(image_path)
# 测试时截取中间的90x90
transformation1 = transforms.Compose([
transforms.CenterCrop(90),
transforms.ToTensor(),
transforms.Normalize([0.5], [0.5])
])
# 预处理图像
image_tensor = transformation1(image).float()
# 额外添加一个批次维度，因为PyTorch将所有的图像当做批次
image_tensor = image_tensor.unsqueeze_(0)
if torch.cuda.is_available():
image_tensor.cuda()
# 将输入变为变量
input = Variable(image_tensor)
# 预测图像的类别
output = model(input)
index = output.data.numpy().argmax()
return index
if __name__ == '__main__':
best_model_path = './output/epoch_462.pth'
model = SimpleNet()
model.load_state_dict(torch.load(best_model_path))
model.eval()
with open('class_names.json', 'r') as f:
class_names = json.load(f)
img_path = './images/test/bubble/066.jpg'
predict_class = class_names[predict_image(model, img_path)]
print(predict_class)
参考链接
- 用pytorch训练图像分类器模型导出ONNX测试
- 用pytorch实现一个简单的图像分类器
