1.首先是写一个nn
from torch import nn
import torch
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.model=nn.Sequential(
nn.Conv2d(3,32,5,1,2),
nn.MaxPool2d(2),
nn.Conv2d(32,32, 5,1,2),
nn.MaxPool2d(2),
nn.Conv2d(32,64,5,1,2),
nn.MaxPool2d(2),
nn.Flatten(),
nn.Linear(64*4*4,64),
nn.Linear(64,10)
)
def forward(self, x):
x = self.model(x)
return x
2.第二部加载数据集,定义模型,进行训练,训练过之后会保存一个权重文件,你可以加载这个权重文件进行再训练,迁移学习。
import torch
import torchvision
import torchvision.transforms as transforms
import torch.nn as nn
import torch.optim as optim
#from model import Net
# 加载CIFAR-10数据集
trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
download=True, transform=transforms.ToTensor())
trainloader = torch.utils.data.DataLoader(trainset, batch_size=128)
testset = torchvision.datasets.CIFAR10(root='./data', train=False,
download=True, transform=transforms.ToTensor())
testloader = torch.utils.data.DataLoader(testset, batch_size=128)
print(len(trainset),len(testset))
# 定义模型
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.model=nn.Sequential(
nn.Conv2d(3,32,5,1,2),
nn.MaxPool2d(2),
nn.Conv2d(32,32, 5,1,2),
nn.MaxPool2d(2),
nn.Conv2d(32,64,5,1,2),
nn.MaxPool2d(2),
nn.Flatten(),
nn.Linear(64*4*4,64),
nn.Linear(64,10)
)
def forward(self, x):
x = self.model(x)
return x
net = Net()
#如果导入了Net,则不需要再写一遍模型
net=torch.load('model_net.pth') #加载训练过的模型,如果没有这一行需要注释
# print(net)
if torch.cuda.is_available():
net=net.cuda()
print('1')
# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
if torch.cuda.is_available():
criterion=criterion.cuda()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
i=0
# 训练模型
for epoch in range(10): # 多次循环遍历数据集
running_loss = 0.0
print(f'----{epoch+1}轮---')
for data in trainloader:
inputs, labels = data
if torch.cuda.is_available():
inputs=inputs.cuda()
labels=labels.cuda()
# print(len(inputs),len(labels))
# 梯度清零
optimizer.zero_grad()
# 前向传播 + 反向传播 + 优化
outputs = net(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
i+=1
print(f'训练第{i}次,loss={loss}')
print('Finished Training')
# 测试模型
correct = 0
total = 0
with torch.no_grad():
for data in testloader:
images, labels = data
if torch.cuda.is_available():
images=images.cuda()
labels=labels.cuda()
outputs = net(images)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
print('Accuracy of the network on the 10000 test images: %d %%' % (
100 * correct / total))
# 保存模型
torch.save(net, 'model_net.pth')
print('Finished Training, model saved as model_net.pth')
3.训练出权重文件,开始进行预测(需要导入Net或者再写一遍模型),out的输出是数据集中的类别序号
import torchvision
import torch
from model import Net
from PIL import Image
image_path='dee.png'
img=Image.open(image_path)
print(img)
r, g, b, a = img.split()
# 创建一个新的RGB图像
rgb_image = Image.merge('RGB', (r, g, b))
transform=torchvision.transforms.Compose(
[torchvision.transforms.Resize((32, 32)),
torchvision.transforms.ToTensor()]
)
img = transform(rgb_image)
print(img.shape)
img = torch.unsqueeze(img, dim=0)
print(img.shape)
model=torch.load('model_net.pth')
print(model)
model.eval()
output=model(img)
print(output)
out=torch.argmax(output)
print(out)