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