pytorch识别图片之003 图片类型判断

import ssl
ssl._create_default_https_context = ssl._create_unverified_context
from torchvision import datasets

from matplotlib import pyplot as plt
import numpy as np
import torch
from torchvision import transforms
data_path = '/Volumes/c/work/aigc'



transforms_cifar10 = datasets.CIFAR10(root=data_path,train=True,download=True,transform=transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.4915,0.4823,0.4468),(0.2470,0.2435,0.2616))]))
import torch.nn as nn
model= nn.Sequential(
    nn.Linear(3*32*32,512),
    nn.Tanh(),
    nn.Linear(512,2), # 输出飞机和鸟的概率 2
    nn.Softmax(dim=1)
)

softmax = nn.Softmax(dim=1) #放大大的 缩小小的
label_map={0:0,2:1}
class_names=['airplane','bird']
cifar2= [(img,label_map[label]) for img,label  in  transforms_cifar10 if label in [0,2]]

img,_=cifar2[0]
plt.imshow(img.permute(1,2,0))
plt.show()
# 数据形状的变化过程
# 原始图片: [3, 32, 32]
#     ↓ view(-1)
# 展平后: [3072]
#     ↓ unsqueeze(0)
# 最终: [1, 3072]  ← 符合模型要求的输入格式
img_batch = img.view(-1).unsqueeze(0)
out = model(img_batch)
print(out)