手写数字识别

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset, TensorDataset
from PIL import Image
import cv2
import torchvision.transforms as transform
import os

path = r'train'


"""
# 提取标签
index = os.listdir(path)
print(index)
for i in index:
    t = i.split('_')
    m = t[2].split('.')[0]
    print(m)
"""

#  处理数据集
class myData(Dataset):
    def __init__(self, path):
        self.path = path
        self.transf = transform.Compose(
            [
                transform.ToTensor(),
                transform.Lambda(lambda x: x.repeat(3, 1, 1)),  #通道扩充为3倍， 长宽为1被
                transform.Normalize(mean=[0.5, 0.5 ,0.5], std=[0.5, 0.5, 0.5])
            ]
        )
        self.path_list = os.listdir(self.path)

    def __getitem__(self, idx):
        img_path = self.path_list[idx]
        abs_img_path = os.path.join(self.path, img_path)
        img = Image.open(abs_img_path)
        img_tensor = self.transf(img)
        # 打标签
        p1 = img_path.split('_')
        label = int(p1[2].split('.')[0])
        label_tensor = torch.as_tensor(label, dtype=torch.int64)

        return img_tensor, label_tensor

    def __len__(self):
        return len(self.path_list)

dataset = myData(path)

trainloader = DataLoader(dataset, batch_size=128, shuffle=True)

# 定义CNN

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1)
        self.conv2 = nn.Conv2d(16, 64, kernel_size=3, stride=1, padding=1)
        self.conv3 = nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1)
        self.relu = nn.ReLU()
        #全连接层
        self.fc1 = nn.Linear(128*28*28, 512)
        self.fc2 = nn.Linear(512, 64)
        self.fc3 = nn.Linear(64, 10)

    def forward(self, x):
        x = self.relu(self.conv1(x))
        x = self.relu(self.conv2(x))
        x = self.relu(self.conv3(x))
        #将x 展开
        x = x.view(-1, 128*28*28)
        x = self.relu(self.fc1(x))
        x = self.relu(self.fc2(x))
        x = self.fc3(x)
        return x

model = CNN()

# 定义损失函数  和  优化器
import torch.optim as optim
from tqdm import tqdm
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

# 记录保存更新 准确率
class AverageMeter(object):
    def __init__(self):
        self.reset()

    def reset(self):
        self.avg = 0
        self.num = 0
        self.cnt = 0

    def update(self, val, n=1):
        self.num += val*n
        self.cnt += n
        self.avg = self.num / self.cnt

# 再多分类问题中， topk准确率是 只要正确标签 的概率 在最大的k个概率中即为正确
def accuracy(output, label, topk=(1,)):
    maxk = max(topk)
    batch_size = label.size(0)
    # get top-k index
    _, pred = output.topk(maxk, 1, True, True)
    pred = pred.t() # transpose
    correct = pred.eq(label.view(1, -1).expand_as(pred))
    # eq方法 逐个比较  相同返回True
    rtn = []
    for k in topk:
        correct_k = correct[:k].contiguous().view(-1).float().sum(0)
        rtn.append(correct_k.mul_(100.0 / batch_size))
    return rtn

device = torch.device("cuda:0"if torch.cuda.is_available() else "cpu")
epochs = 10
#model = model.to(device)


"""
#训练
for epoch in range(epochs):
    model.train()
    train_loss = 0.0
    ac = AverageMeter()
    trainloader = tqdm(trainloader)
    trainloader.set_description('[%s:%04d/%04d]' % ('Epoch', epoch+1, epochs))
    for i , data in enumerate(trainloader):
        x, labels = data[0].to(device), data[1].to(device)
        optimizer.zero_grad()
        pred = model(x)
        batch_loss = criterion(pred, labels)
        batch_loss.backward()
        optimizer.step()
        #计算准确率
        acc1, acc2 = accuracy(pred, labels, topk=(1, 3))
        n = x.size(0)
        ac.update(acc1, n)
        train_loss += batch_loss
        postfix = {"train_loos: %.6f" % (train_loss/(i+1)), "train_acc: %6f" % (ac.avg)}
        trainloader.set_postfix(log=postfix)

torch.save(model.state_dict(), 'model_resnet.pth')
print("训练完成")

"""
image_path = r'test/test_147_2.jpg'
img = Image.open(image_path)
image = transform.ToTensor()(img)
image = transform.Lambda(lambda x: x.repeat(3, 1, 1))(image)
image = transform.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))(image).to(device)

state_dict = torch.load('model.pth')
mo = CNN()
mo.to(device)
mo.load_state_dict(state_dict)

out = mo(image)

softmax = nn.Softmax(dim=1)
res = softmax(out)
print(res)
print("图片的数字是： ", torch.argmax(res).item())

ii = cv2.imread(image_path)
cv2.imshow('image', ii)
cv2.waitKey(0)
cv2.destroyAllWindows()

结果

C:\Users\16630\Desktop\file\cc\venv\Scripts\python.exe C:/权当作D盘/临时文件/WritngNum/shoxuie.py 
tensor([[2.1309e-17, 3.6944e-12, 1.0000e+00, 1.0493e-13, 3.4431e-15, 1.4164e-18,
         3.1761e-18, 3.5542e-14, 2.5836e-14, 8.1047e-17]], device='cuda:0',
       grad_fn=<SoftmaxBackward0>)
图片的数字是：  2

2 对应 位置上的概率几乎为1

C:\Users\16630\Desktop\file\cc\venv\Scripts\python.exe C:/权当作D盘/临时文件/WritngNum/shoxuie.py 
tensor([[1.5335e-18, 5.0806e-18, 9.4139e-13, 4.2161e-11, 1.1821e-19, 6.5505e-14,
         1.0558e-17, 2.0131e-16, 1.0000e+00, 2.4899e-15]], device='cuda:0',
       grad_fn=<SoftmaxBackward0>)
图片的数字是：  8