from torch import nn
from torch.nn.functional import relu
class BaseBlock(nn.Module):
def __init__(self, in_channels, out_channels, stride):
super(BaseBlock, self).__init__()
self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1)
self.bn1 = nn.BatchNorm2d(out_channels)
self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=stride, padding=1)
self.bn2 = nn.BatchNorm2d(out_channels)
def forward(self, x):
out_put = self.conv1(x)
out_put = relu(self.bn1(out_put))
out_put = self.conv2(out_put)
out_put = self.bn2(out_put)
out_put = relu(out_put + out_put)
return out_put
class DownBlock(nn.Module):
def __init__(self, in_channels, out_channels, stride):
super(DownBlock, self).__init__()
self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride[0], padding=1)
self.bn1 = nn.BatchNorm2d(out_channels)
self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=stride[1], padding=1)
self.bn2 = nn.BatchNorm2d(out_channels)
self.extra = nn.Sequential(
nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride[0], padding=0),
nn.BatchNorm2d(out_channels)
)
def forward(self, x):
extra_x = self.extra(x)
out_put = self.conv1(x)
out_put = relu(self.bn1(out_put))
out_put = self.conv2(out_put)
out_put = self.bn1(out_put)
out_put = relu(out_put + extra_x)
return out_put
from torch import nn
from torch.nn.functional import relu
from .base_block import BaseBlock, DownBlock
class RestNEt18(nn.Module):
def __init__(self, in_channels):
super(RestNEt18, self).__init__()
self.conv1 = nn.Conv2d(1, 64, 7, 2, 3)
self.bn1 = nn.BatchNorm2d(64)
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
self.layer1 = nn.Sequential(BaseBlock(64, 64, 1),
BaseBlock(64, 64, 1))
self.layer2 = nn.Sequential(DownBlock(64, 128, [2, 1]),
BaseBlock(128, 128, 1))
self.layer3 = nn.Sequential(DownBlock(128, 256, [2, 1]),
BaseBlock(256, 256, 1))
self.layer4 = nn.Sequential(DownBlock(256, 512, [2, 1]),
BaseBlock(512, 512, 1))
# 二位自适应平均池化，根据输入的尺寸以及设定的输出大小计算出输出元素在输入中的感受野，然后再感受野上进行平均池化
self.agv_pool = nn.AdaptiveAvgPool2d((1, 1))
self.fc = nn.Linear(512, 10)
def forward(self, x):
out_put = self.conv1(x)
# print(out_put.shape)
out_put = relu(self.bn1(out_put))
out_put = self.maxpool(out_put)
# print(out_put.shape)
out_put = self.layer1(out_put)
# print(out_put.shape)
out_put = self.layer2(out_put)
# print(out_put.shape)
out_put = self.layer3(out_put)
# print(out_put.shape)
out_put = self.layer4(out_put)
# print(out_put.shape)
out_put = self.agv_pool(out_put)
# print(out_put.shape)
out_put = out_put.reshape(out_put.shape[0], -1)
# print(out_put.shape)
out_put = self.fc(out_put)
# print(out_put.shape)
return out_put
from abc import ABCMeta, abstractmethod
class TrainBase(metaclass=ABCMeta):
@abstractmethod
def _get_data_loader(self, is_train=True):
"""
获取训练数据和测试数据
:param is_train: 标识测试数据或训练数据
:return: (Feature, target)
"""
pass
@abstractmethod
def _train(self, my_model, loss_fn, optimizer, train_data):
"""
训练模型
:param my_model: 模型
:param loss_fn: 损失函数
:param optimizer: 优化器
:param train_data: 训练数据
:return:
"""
@abstractmethod
def _test(self, my_model, test_data):
"""
测试模型
:param my_model: 模型
:param test_data: 测试数据
:return:
"""
pass
@abstractmethod
def predict(self, img_path):
"""
使用模型预测数据
:param img_path: 图片文件地址
:return: 预测值
"""
pass
def run(self):
"""
执行深度学习模型训练
:return:
"""
pass
import os.path
from torchvision import transforms
from torchvision import datasets
from torch import nn, optim
from torch.utils.tensorboard import SummaryWriter
from torch.utils.data import DataLoader
from PIL import Image
from models.resnet.resnet18 import RestNEt18
from Itrain import TrainBase
import numpy as np
import torch
class NumPredict(TrainBase):
def __init__(self, model_path, epoch=100, lr=0.001):
# 定义超参数
self.epoch = epoch
self.lr = lr
self.model_path = model_path
# 固定参数
self.current_train_step = 0
self.acc = 0
self.test_count = 0
self.writer = None
self.transforms = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(mean=0.1307, std=0.3081)
])
def _get_data_loader(self, is_train=True):
data = datasets.MNIST(root="../datasets", train=is_train, transform=self.transforms, download=True)
return data
def _train(self, my_model, loss_fn, optimizer, train_data):
# 训练模型
for data in train_data:
imgs, target = data
# 计算预测值
predict = my_model(imgs)
# 计算损失值
loss = loss_fn(predict, target)
# 梯度归0
optimizer.zero_grad()
# 反向计算梯度
loss.backward()
# 反向更新
optimizer.step()
# 输出损失值
self.current_train_step += 1
if self.current_train_step % 10 == 0:
print("训练次数%s 损失值：%s" % (self.current_train_step, loss.item()))
self.writer.add_scalar("train_loss", loss.item(), global_step=self.current_train_step)
def _test(self, my_model, test_data):
# 测试模型
total_acc = 0
for data in test_data:
imgs, target = data
with torch.no_grad():
predict = my_model(imgs)
acc = (predict.argmax(1) == target).sum()
print("准确率：%s" % str(acc/imgs.shape[0]))
total_acc += acc
self.writer.add_scalar("test_acc", total_acc / self.test_count, global_step=self.current_train_step)
def run(self):
"""
:return:
"""
# 准备数据
train_data = self._get_data_loader()
train_data = DataLoader(train_data, batch_size=64, shuffle=True)
test_data = self._get_data_loader(is_train=False)
self.test_count = len(test_data)
test_data = DataLoader(test_data, batch_size=1000, shuffle=True)
# 定义神经网络
my_model = RestNEt18(in_channels=1)
# 断点续训
if os.path.exists(self.model_path):
my_model = torch.load(self.model_path)
# 定义损失函数
loss_fn = nn.CrossEntropyLoss()
# 定义优化器
optimizer = optim.SGD(my_model.parameters(), lr=self.lr)
self.writer = SummaryWriter(log_dir="../logs")
for item in range(self.epoch):
print("-----开始第%s轮测试-----" % str(item))
# 训练模型
my_model.train()
self._train(my_model, loss_fn, optimizer, train_data)
# 测试模型
my_model.eval()
self._test(my_model, test_data)
# 存储模型
torch.save(my_model, self.model_path)
def predict(self, img_path):
"""
使用模型预测数据
:param img_path: 图片文件地址
:return: 预测值
"""
# 读取图片
img = Image.open(img_path)
# 图片预处理满足模型需求
# png图片是四通道转化为3通道
img = img.convert("RGB")
# 转为灰度图
img = img.convert("L")
# 二值化
img = img.point(lambda x: 255 if x >= 100 else 0)
# img.show()
# 由于训练数据以黑色为底的图片，预测图片为白色为底，需要进行转换
img = img.point(lambda x: 255-x)
# img.show()
# 图片尺寸修改，转化为张量
img = self.transforms(img)
# 加载模型
my_model = torch.load(self.model_path)
# 预测
img = torch.reshape(img, (1, 1, 224, 224))
my_model.eval()
with torch.no_grad():
predict = my_model(img)
print(predict)
predict = predict.argmax(1)
print(predict.item())
return predict
if __name__ == "__main__":
# 生成模型
# pth_path = "../pth/num_predict2.pth"
# num_predict = NumPredict(pth_path, epoch=1)
# num_predict.run()
# 预测数据
img_path = "../img/9.png"
pth_path = "../pth/num_predict2.pth"
num_predict = NumPredict(pth_path)
num_predict.predict(img_path)