# -*- coding: UTF-8 -*-
# Created by -牧野- CSDN https://blog.csdn.net/dcrmg/article/details/102807575
# 参考 https://blog.csdn.net/JR_Chan/article/details/95641758
import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
import time
import os
from math import ceil
import argparse
import copy
from PIL import Image
from torchvision import transforms, datasets
from torch.autograd import Variable
from tensorboardX import SummaryWriter
# 定义一个简单的二分类网络
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
# 训练模型入口
def train(args):
# read data
dataloders, dataset_sizes, class_names = ImageDataset(args)
with open(args.class_file, 'w') as f:
for name in class_names:
f.writelines(name + '\n')
# 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=1e-4)
# 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,
dataloders = dataloders)
torch.save(model.state_dict(), os.path.join(args.save_path, 'best_model.pth'))
writer.close()
# 测试单张图片(使用pth模型)入口
def test(test_model_path, test_img_path, class_file):
best_model_path = test_model_path
model = SimpleNet()
model.load_state_dict(torch.load(best_model_path))
model.eval()
class_names = []
with open(class_file, 'r') as f:
lines = f.readlines()
for line in lines:
class_names.append(line)
img_path = test_img_path
predict_class = class_names[predict_image(model, img_path)]
print(predict_class)
# 转换pytorch训练的pth模型到ONNX模型
def convert_model_to_ONNX(input_img_size, input_pth_model, output_ONNX):
dummy_input = torch.randn(3, 1, input_img_size, input_img_size)
model = SimpleNet()
state_dict = torch.load(input_pth_model, map_location='cpu')
model.load_state_dict(state_dict)
model.eval() # 设置模型为推理模式(重要)
input_names = ["input_image"]
output_names = ["output_classification"]
torch.onnx.export(model, dummy_input, output_ONNX, verbose=True, input_names=input_names,
output_names=output_names)
# 训练模型主函数
def train_model(args, model, criterion, optimizer, scheduler, num_epochs, dataset_sizes, use_gpu, dataloders):
begin = 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 and os.path.exists('result.txt'):
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() - begin
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
# 测试单张图片主函数
def predict_image(model, image_path):
image = Image.open(image_path).convert('L')
# 测试时截取中间的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
# 使用PIL读取图片并转换为灰度图
def readImg(path):
im = Image.open(path)
return im.convert("L")
# 读取训练和测试数据
def ImageDataset(args):
# 数据增强及归一化
# 图片都是120x120的,训练时随机裁取90x90的部分,测试时裁取中间的90x90
data_transforms = {
'train': transforms.Compose([
transforms.RandomCrop(def_img_train_and_test_size),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize([0.5], [0.5])
]),
'test': transforms.Compose([
transforms.CenterCrop(def_img_train_and_test_size),
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
# 设置参数
def set_parser():
parser = argparse.ArgumentParser(description='classification')
# 图片数据的根目录(Root catalog of images)
parser.add_argument('--data-dir', type=str, default='images')
parser.add_argument('--class-file', type=str, default='class_names.class')
parser.add_argument('--batch-size', type=int, default=8)
parser.add_argument('--num-epochs', type=int, default=100)
parser.add_argument('--lr', type=float, default=0.002) # those who set lr greater than 0.01 are hooligans!!
parser.add_argument('--num-workers', type=int, default=8)
parser.add_argument('--print-freq', type=int, default=100)
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')
return parser.parse_args()
if __name__ == '__main__':
writer = SummaryWriter(log_dir='log')
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
def_img_train_and_test_size = 90 # 训练尺寸
args = set_parser() # 设置参数
train(args) # 训练模型
test('./output/best_model.pth', './images/test/cat/cat.0.jpg', args.class_file) # 测试模型(单张图片)
# 转换pytorch的pth模型到ONNX模型
convert_model_to_ONNX(def_img_train_and_test_size, './output/epoch_99.pth', "./cat_dog_classify.onnx")
onnx.cpp
// PthONNX.cpp : 基于OpenCV dnn、 onnx 的cat、dog二分类程序
// Created by -牧野- 2019年10月29日 https://blog.csdn.net/dcrmg/article/details/102807575
//
#include <iostream>
#include <opencv2/highgui.hpp>
#include <opencv2/imgproc.hpp>
#include <opencv2/dnn.hpp>
#include <fstream>
//ONNX 执行推理类
class PthONNX {
public:
//@model_path ONNX模型路径
//@classes_file_path 分类信息文件
//@input_size 网络输入大小
PthONNX(const std::string &model_path, const std::string &classes_file_path, cv::Size input_size);
//@input_image 输入图片,BGR格式
//@classification_output 网络输出的分类名称 0:cat 1:dog 1:None
void Classify(const cv::Mat &input_image, std::string &classification_output);
private:
void ClassifyImplement(const cv::Mat &image, std::string &classification_output);
private:
cv::Size input_size_;
cv::dnn::Net net_classify_;
std::vector<std::string> classes_;
};
// 构造函数
PthONNX::PthONNX(const std::string &model_path, const std::string &classes_file_path,
cv::Size input_size) : input_size_(input_size) {
std::ifstream ifs(classes_file_path.c_str());
assert(ifs.is_open());
std::string line;
while (getline(ifs, line)) {
line = line;
classes_.push_back(line);
}
net_classify_ = cv::dnn::readNetFromONNX(model_path);
net_classify_.setPreferableBackend(cv::dnn::DNN_BACKEND_OPENCV);
net_classify_.setPreferableTarget(cv::dnn::DNN_TARGET_CPU);
}
// ONNX推理入口函数
void PthONNX::Classify(const cv::Mat &input_image, std::string &classification_results) {
assert(input_image.data);
cv::Mat image = input_image.clone();
cv::resize(image, image, cv::Size(90, 90));
cv::cvtColor(image, image, cv::COLOR_BGR2GRAY);
ClassifyImplement(image,classification_results);
}
//ONNX推理主函数
void PthONNX::ClassifyImplement(const cv::Mat &image,std::string &classification_results) {
classification_results.clear();
//***********前处理***********
cv::Scalar mean_value(0, 0, 0);
cv::Mat input_blob = cv::dnn::blobFromImage(image, 1, input_size_, mean_value, false, false, CV_32F);
//***********前处理***********
net_classify_.setInput(input_blob);
const std::vector<cv::String> &out_names = net_classify_.getUnconnectedOutLayersNames();
cv::Mat out_tensor = net_classify_.forward(out_names[0]);
//***********后处理***********
double minVal;
double maxVal;
cv::Point minIdx;
cv::Point maxIdx; // minnimum Index, maximum Index
cv::minMaxLoc(out_tensor, &minVal, &maxVal, &minIdx, &maxIdx);
int index_class = maxIdx.x;
classification_results = (index_class <= 1) ? classes_[index_class] : "None";
//***********后处理***********
}
int main()
{
const std::string img_path = "D:/1/1/SimpleNet-master/images/train/cat/cat.4896.jpg";
const std::string onnx_model_path = "D:/1/1/pytorch-train-test-onnx/cat_dog_classify.onnx";
const std::string class_names_file_path = "D:/software/VS2019_Test/PthONNX/x64/class_names.class";
const cv::Size net_input_size(90, 90);
cv::Mat img = cv::imread(img_path);
std::string classify_output; // 分类结果
PthONNX classifier(onnx_model_path, class_names_file_path, net_input_size);
classifier.Classify(img, classify_output);
std::cout << "图片类别:" << classify_output << std::endl << std::endl;
cv::putText(img, classify_output, cv::Point(20,20), 2, 1.2, cv::Scalar(0, 0, 255));
cv::imshow("classify", img);
cv::waitKey();
}