#_*_ coding:utf-8 _*_
# pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu117
import torch
from torch import nn
from torch.nn import Conv2d,MaxPool2d,Linear,Sequential,Flatten
from torch.utils.data import DataLoader
import torch.optim.optimizer
import torchvision
from torch.utils.tensorboard import SummaryWriter
print(SummaryWriter)
class Model(nn.Module):
def __init__(self):
super().__init__()
self.covn1=Conv2d(in_channels=3,out_channels=32,kernel_size=5,padding=2)
self.maxpool1=MaxPool2d(kernel_size=2)
self.covn2=Conv2d(in_channels=32,out_channels=32,kernel_size=5,padding=2)
self.maxpool2=MaxPool2d(kernel_size=2)
self.covn3=Conv2d(in_channels=32,out_channels=64,kernel_size=5,padding=2)
self.maxpool3=MaxPool2d(kernel_size=2)
self.flaten=Flatten()
self.linear1=Linear(in_features=1024,out_features=64)
self.linear2=Linear(in_features=64,out_features=10)
def forward(self,x):
x= self.covn1(x)
x=self.maxpool1(x)
x=self.covn2(x)
x=self.maxpool2(x)
x=self.covn3(x)
x=self.maxpool3(x)
x=self.flaten(x)
x=self.linear1(x)
x=self.linear2(x)
return x
if __name__ == "__main__":
model = Model()
print(model)
#下载数据集
train_data = torchvision.datasets.CIFAR10('../CIFAR-10/',train=True,transform=torchvision.transforms.ToTensor(),download=False)
test_data = torchvision.datasets.CIFAR10('../CIFAR-10/',train=False,transform=torchvision.transforms.ToTensor(),download=False)
#数据加载器
train_data_loader = DataLoader(train_data,batch_size=64)
test_data_loader = DataLoader(test_data,batch_size=64)
#损失函数
loss = nn.CrossEntropyLoss()
#优化器
optim = torch.optim.SGD(model.parameters(),lr=0.01)
#创建可视化
write = SummaryWriter('CIFAR_logs')
#开始训练数据
eplo =10
train_step = 0
for i in range(10):
print("开始训练:".format(i))
for img,laber in train_data_loader:
input_d = model(img)
loss_fn =loss(input_d,laber)
optim.zero_grad() #梯度清零
loss_fn.backward() #反向计算
optim.step()#更新梯度
train_step += 1
if train_step%100==0:
print("训练次数：{}".format(train_step),"损失：{}".format(loss_fn.item()))
write.add_scalar('train_loss',loss_fn.item(),train_step)
total_accuracy = 0
total_test_loss = 0
total_test_step = 0
with torch.no_grad():
for imgs,labler in test_data_loader:
output = model(imgs)
loss_re=loss(output,labler)
accuracy = (output.argmax(1) == labler).sum()
total_accuracy += accuracy
total_test_loss += loss_re
total_test_step += 1
print("测试集上的正确率{}".format(total_accuracy/len(test_data)))
print("测试集上的损失{}".format(total_test_loss))
write.add_scalar("test_loss",total_test_loss.item(),total_test_step)
torch.save(model,"model10")
write.close()
# tensorboard --logdir=CIFAR_logs --port=2017预测代码：
from torchvision import datasets, transforms
import numpy as np
from PIL import Image
import torch
import torch.nn.functional as F
from cov01 import Model
classes = ('plane', 'car', 'bird', 'cat',
'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
if __name__ == '__main__':
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = torch.load('model10') # 加载模型
model = model.to(device)
model.eval() # 把模型转为test模式
img = Image.open("../bird.jpg")
trans = transforms.Compose(
[
transforms.CenterCrop(32),
transforms.ToTensor(),
# transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
])
img = trans(img)
img = img.to(device)
img = img.unsqueeze(0) # 图片扩展多一维,因为输入到保存的模型中是4维的[batch_size,通道,长，宽]，而普通图片只有三维，[通道,长，宽]
# 扩展后，为[1，1，28，28]
output = model(img)
prob = F.softmax(output, dim=1) # prob是10个分类的概率
print(prob)
value, predicted = torch.max(output.data, 1) #按照维度返回最大概率dim = 0 表示按列求最大值，并返回最大值的索引，dim = 1 表示按行求最大值，并返回最大值的索引
print(predicted.item())
print(value)
pred_class = classes[predicted.item()]
print(pred_class)
'''
记住：
torch.max()[0]， 只返回最大值的每个数
troch.max()[1]， 只返回最大值的每个索引
torch.max()[1].data 只返回variable中的数据部分（去掉Variable containing:）
torch.max()[1].data.numpy() 把数据转化成numpy ndarry
torch.max()[1].data.numpy().squeeze() 把数据条目中维度为1 的删除掉
torch.max(tensor1,tensor2) element-wise 比较tensor1 和tensor2 中的元素，返回较大的那个值
'''成功预测为bird