重写Dataset
没有用官方导入的数据集,而是自己读取图片和标签,主要熟悉一下处理数据集的流程。
dataset.py
from torch.utils.data import Dataset
from PIL import Image # 读取图片
import os
# 创建一个class,继承Dataset类
class MyDataset(Dataset):
def __init__(self, path, transform=None):
self.path = path
self.transform = transform
self.image_path = os.path.join(self.path, 'images') # 获得图片的路径
self.label_path = os.path.join(self.path, 'labels') # 获得标签的路径
self.image_list = os.listdir(self.image_path) # 获得图片列表
self.label_list = os.listdir(self.label_path) # 获得标签列表
def __getitem__(self, idx): # idx为编号
# 获取图片
img_name = self.image_list[idx] # 名称
img_item_path = os.path.join(self.image_path, img_name)
img = Image.open(img_item_path).convert('RGB') # 读取图片
if self.transform is not None:
img = self.transform(img)
# 获取标签
label_name = self.label_list[idx] # 名称
label_item_path = os.path.join(self.label_path, label_name)
with open(label_item_path, 'r') as f:
label = int(f.read())
return img, label
def __len__(self): # 数据集的长度
return len(self.image_list)
# # 用类创建实例
# path = 'data/cifar10/train/'
# dataset = MyDataset(path)
# img, label = dataset[0]
# img.show()
# print(label)
# a = 1
Model
model.py
# model.py
import torch.nn as nn
import torch.nn.functional as F
class LeNet(nn.Module):
def __init__(self):
super(LeNet, self).__init__()
self.conv1 = nn.Conv2d(3, 16, 5)
self.pool1 = nn.MaxPool2d(2, 2)
self.conv2 = nn.Conv2d(16, 32, 5)
self.pool2 = nn.MaxPool2d(2, 2)
self.fc1 = nn.Linear(32 * 5 * 5, 120)
self.fc2 = nn.Linear(120, 84)
self.fc3 = nn.Linear(84, 10)
def forward(self, x):
x = F.relu(self.conv1(x)) # input(3,32,32) output(16,28,28)
x = self.pool1(x) # output(16,14,14)
x = F.relu(self.conv2(x)) # output(32,10.10)
x = self.pool2(x) # output(32,5,5)
x = x.view(-1, 32 * 5 * 5) # output(5*5*32)
x = F.relu(self.fc1(x)) # output(120)
x = F.relu(self.fc2(x)) # output(84)
x = self.fc3(x) # output(10)
return x
# #model调试
# import torch
#
# #定义shape
# input1 = torch.rand([36,3,32,32])
# print(input1)
# model = LeNet()#实例化
# print(model)
# #输入网络中
# output = model(input1)
# print(output)
# a = 1
Train
train.py
# train.py
import torch
import torchvision
import torch.nn as nn
from dataset import MyDataset
from model import LeNet
import torch.optim as optim
import torchvision.transforms as transforms
from torchvision import transforms, datasets, utils
import matplotlib.pyplot as plt
import numpy as np
# device : GPU or CPU
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(device)
# 图像变换
transform = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
# 训练图片
train_path = 'data/cifar10/train/'
train_set = MyDataset(train_path, transform=transform)
train_loader = torch.utils.data.DataLoader(train_set, batch_size=36,
shuffle=True, num_workers=0)
# 验证图片
val_set = MyDataset(train_path, transform=transform)
val_loader = torch.utils.data.DataLoader(val_set, batch_size=50,
shuffle=True, num_workers=0)
val_data_iter = iter(val_loader)
val_image, val_label = val_data_iter.next()
print(val_image.size())
# 网络
net = LeNet()
net.to(device)
# 损失函数
loss_function = nn.CrossEntropyLoss()
# 定义优化器
optimizer = optim.Adam(net.parameters(), lr=0.001)
# 训练过程
for epoch in range(10): # loop over the dataset multiple times
running_loss = 0.0 # 累加损失
for step, data in enumerate(train_loader, start=0):
# get the inputs; data is a list of [inputs, labels]
inputs, labels = data
# print(inputs.size(), labels.size())
# zero the parameter gradients
optimizer.zero_grad() # 如果不清除历史梯度,就会对计算的历史梯度进行累加
# forward + backward + optimize
outputs = net(inputs.to(device))
loss = loss_function(outputs, labels.to(device))
loss.backward()
optimizer.step()
# print statistics
running_loss += loss.item()
if step % 500 == 499: # print every 500 mini-batches
with torch.no_grad(): # 上下文管理器
outputs = net(val_image.to(device)) # [batch, 10]
predict_y = torch.max(outputs, dim=1)[1]
accuracy = (predict_y == val_label.to(device)).sum().item() / val_label.size(0)
print('[%d, %5d] train_loss: %.3f test_accuracy: %.3f' %
(epoch + 1, step + 1, running_loss / 500, accuracy))
running_loss = 0.0
print('Finished Training')
# 保存模型
save_path = './Lenet.pth'
torch.save(net.state_dict(), save_path)
训练过程
Test
test.py
# test.py
import torch
import torchvision.transforms as transforms
from PIL import Image
from dataset import MyDataset
from model import LeNet
import matplotlib.pyplot as plt
import numpy as np
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# 预处理
transform = transforms.Compose(
[transforms.Resize((32, 32)),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
# 测试图片
test_path = 'data/cifar10/test/'
test_set = MyDataset(test_path, transform=transform)
test_loader = torch.utils.data.DataLoader(test_set, batch_size=36,
shuffle=True, num_workers=0)
classes = ('plane', 'car', 'bird', 'cat',
'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
net = LeNet()
net.load_state_dict(torch.load('./Lenet.pth'))
net.to(device)
correct = 0
total = 0
for step, data in enumerate(test_loader, start=0):
with torch.no_grad():
inputs, labels = data
inputs = inputs.to(device)
labels = labels.to(device)
outputs = net(inputs.to(device))
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
print('Test Accuracy of the model on the 10000 test images: {} %'.format(100 * correct / total))
测试集结果