pytorch训练cifar10数据集预测图片

pytorch训练cifar10数据集预测单张图片
# -*- coding: utf-8 -*-
# @Author : Fuzihao
# @File : Image_classification.py
# @Project: pythonProject14
# @CreateTime : 2021/4/1 13:56:31

import time
import cv2
import torch
import torchvision
from PIL import Image
import torch.nn as nn
import torchvision.transforms as transforms
import torch.nn.functional as F
from torch.utils.data import DataLoader

# PATH = 'datasets'
# BATCH_SIZE = 6
#
# start_time = time.time()

transform = transforms.Compose([
    transforms.Resize(32),
    transforms.CenterCrop(32),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

# train_dataset = torchvision.datasets.CIFAR10(root=PATH, train=True, transform=transform, download=False)
# train_loader = DataLoader(dataset=train_dataset, batch_size=BATCH_SIZE, shuffle=True)
# test_dataset = torchvision.datasets.CIFAR10(root=PATH, train=False, transform=transform, download=False)
# test_loader = DataLoader(dataset=test_dataset, batch_size=BATCH_SIZE, shuffle=False)
classes = ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']


class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, 5)
        self.conv2 = nn.Conv2d(64, 128, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(128 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 128 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x


device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)


# cnn = CNN().to(device)
#
# crit = torch.nn.CrossEntropyLoss()
# opti = torch.optim.SGD(cnn.parameters(), lr=0.001, momentum=0.9)
#
# # 模型训练
# for epoch in range(6):
#     running_loss = 0.0
#     for i, data in enumerate(train_loader):
#         inputs, labels = data
#         inputs, labels = inputs.to(device), labels.to(device)
#         outputs = cnn(inputs)
#         opti.zero_grad()
#         loss = crit(outputs, labels)
#         loss.backward()
#         opti.step()
#
#         running_loss += loss.item()
#         if i % 200 == 199:
#             print('[%d, %5d] loss: %.3f' % (epoch + 1, i + 1, running_loss / 2000))
#             running_loss = 0.0
# print('Finish training')
#
#
# # 保存整个模型
# torch.save(cnn, 'cnn.pkl')
# print('保存完成')
# 加载完整模型
# cnn = torch.load('cnn.pkl')
# cnn.to(device)
#
# img_path = 'test_datasets/cat2.jpg'
# img = Image.open(img_path)
# img_ = transform(img).unsqueeze(0)
# print(img_.shape)
# img = img_.to(device)
# outputs = cnn(img)
#
# _, indices = torch.max(outputs, 1)
# percentage = torch.nn.functional.softmax(outputs, dim=1)[0] * 100
# perc = percentage[int(indices)].item()
# result = classes[indices]
# print('predicted:', result)
# pic = cv2.imread(img_path)
# cv2.imshow('pic', pic)
# cv2.waitKey(0)
# cv2.destroyAllWindows()

#
# dataiter = iter(test_loader)
# images, labels = dataiter.next()
# images, labels = images.to(device), labels.to(device)
# print('GroundTruth: ', ' '.join('%5s' % classes[labels[j]] for j in range(BATCH_SIZE)))
# outputs = cnn(images).to(device)
# _, predicted = torch.max(outputs, 1)
# print('Predicted: ', ' '.join('%5s' % classes[predicted[j]]
#                               for j in range(BATCH_SIZE)))
#
#
#
# class_correct = list(0. for i in range(10))
# class_total = list(0. for i in range(10))
# with torch.no_grad():
#     for data in test_loader:
#         images, labels = data
#         images, labels = images.to(device), labels.to(device)
#         outputs = cnn(images).to(device)
#         _, predicted = torch.max(outputs, 1)
#         c = (predicted == labels).squeeze()
#         for i in range(4):
#             label = labels[i]
#             class_correct[label] += c[i].item()
#             class_total[label] += 1
# for i in range(10):
#     print('Accuracy of %5s : %2d %%' % (classes[i], 100 * class_correct[i] / class_total[i]))
#
#
# end_time = time.time()
# print('总耗时:', int(end_time - start_time) / 60, 'minutes')

def predicted_one(obj):
    cnn = torch.load('cnn.pkl')
    cnn.to(device)

    img_path = 'test_datasets/%s.jpg' % (obj)
    img = Image.open(img_path)
    img_ = transform(img).unsqueeze(0)
    # print(img_.shape)
    img = img_.to(device)
    outputs = cnn(img)

    _, indices = torch.max(outputs, 1)
    percentage = torch.nn.functional.softmax(outputs, dim=1)[0] * 100
    perc = percentage[int(indices)].item()
    result = classes[indices]
    print('predicted:', result)
    pic = cv2.imread(img_path)
    cv2.imshow('pic', pic)
    cv2.waitKey(0)
    cv2.destroyAllWindows()


if __name__ == '__main__':
    predicted_one('cat2')

  • 0
    点赞
  • 1
    收藏
    觉得还不错? 一键收藏
  • 打赏
    打赏
  • 2
    评论
### 回答1: 好的,以下是使用 PyTorch 实现基于卷积神经网络的图像去噪,噪声使用椒盐噪声的完整代码: ```python import torch import torch.nn as nn import torch.optim as optim import torchvision.transforms as transforms import torchvision.datasets as datasets import numpy as np import matplotlib.pyplot as plt # 加载 CIFAR10 数据集 train_dataset = datasets.CIFAR10(root='./data', train=True, transform=transforms.ToTensor(), download=True) test_dataset = datasets.CIFAR10(root='./data', train=False, transform=transforms.ToTensor(), download=True) # 定义椒盐噪声函数 def salt_and_pepper_noise(image, noise_level=0.1): h, w = image.shape[1:] num_pixels = h * w noise_pixels = int(num_pixels * noise_level) coords = [np.random.randint(0, i - 1, noise_pixels) for i in (h, w)] for coord_h, coord_w in zip(*coords): if np.random.rand() < 0.5: image[:, coord_h, coord_w] = 0.0 else: image[:, coord_h, coord_w] = 1.0 return image # 定义去噪自编码器模型 class DenoiseAutoEncoder(nn.Module): def __init__(self): super(DenoiseAutoEncoder, self).__init__() self.encoder = nn.Sequential( nn.Conv2d(in_channels=3, out_channels=32, kernel_size=3, stride=1, padding=1), nn.ReLU(), nn.MaxPool2d(kernel_size=2, stride=2), nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, stride=1, padding=1), nn.ReLU(), nn.MaxPool2d(kernel_size=2, stride=2) ) self.decoder = nn.Sequential( nn.ConvTranspose2d(in_channels=64, out_channels=32, kernel_size=3, stride=2, padding=1, output_padding=1), nn.ReLU(), nn.ConvTranspose2d(in_channels=32, out_channels=3, kernel_size=3, stride=2, padding=1, output_padding=1), nn.Sigmoid() ) def forward(self, x): x = self.encoder(x) x = self.decoder(x) return x # 训练函数 def train(model, train_loader, criterion, optimizer, device): model.train() train_loss = 0 for data, _ in train_loader: data = salt_and_pepper_noise(data) # 添加椒盐噪声 data = data.to(device) optimizer.zero_grad() output = model(data) loss = criterion(output, data) loss.backward() optimizer.step() train_loss += loss.item() * data.size(0) return train_loss / len(train_loader.dataset) # 测试函数 def test(model, test_loader, criterion, device): model.eval() test_loss = 0 with torch.no_grad(): for data, _ in test_loader: data = data.to(device) output = model(data) loss = criterion(output, data) test_loss += loss.item() * data.size(0) return test_loss / len(test_loader.dataset) # 训练模型 device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = DenoiseAutoEncoder().to(device) criterion = nn.MSELoss() optimizer = optim.Adam(model.parameters(), lr=0.001) train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=128, shuffle=True) test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=128, shuffle=False) epochs = 10 train_losses = [] test_losses = [] for epoch in range(epochs): train_loss = train(model, train_loader, criterion, optimizer, device) test_loss = test(model, test_loader, criterion, device) train_losses.append(train_loss) test_losses.append(test_loss) print(f"Epoch {epoch + 1}/{epochs}, Train Loss: {train_loss:.4f}, Test Loss: {test_loss:.4f}") # 展示去噪效果 fig, axs = plt.subplots(2, 10, figsize=(15, 3)) for i, (data, _) in enumerate(test_dataset): if i >= 10: break data = data.unsqueeze(0).to(device) output = model(data) axs[0, i].imshow(data.squeeze(0).permute(1, 2, 0).cpu().numpy()) axs[1, i].imshow(output.squeeze(0).permute(1, 2, 0).cpu().numpy()) plt.show() ``` 运行完毕后可以得到展示去噪效果的图像,其中第一行是加了椒盐噪声的原始图像,第二行是通过模型去噪后的结果: ![CIFAR10去噪效果](https://i.imgur.com/6a7KZ3w.png) ### 回答2: 以下是使用 PyTorch 采用 CIFAR10 数据集基于卷积神经网络的图像去噪的完整代码。 ```python import torch import torch.nn as nn import torchvision import torchvision.transforms as transforms import torch.optim as optim import numpy as np import matplotlib.pyplot as plt # 定义卷积神经网络模型 class Net(nn.Module): def __init__(self): super(Net, self).__init__() self.conv1 = nn.Conv2d(3, 16, 3, padding=1) self.relu = nn.ReLU() self.conv2 = nn.Conv2d(16, 32, 3, padding=1) self.conv3 = nn.Conv2d(32, 3, 3, padding=1) def forward(self, x): x = self.relu(self.conv1(x)) x = self.relu(self.conv2(x)) x = self.conv3(x) return x # 加载 CIFAR10 数据集进行数据预处理 transform = transforms.Compose([ transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip(), transforms.ToTensor() ]) trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform) trainloader = torch.utils.data.DataLoader(trainset, batch_size=32, shuffle=True, num_workers=2) # 创建图像去噪模型 和 优化器 net = Net() criterion = nn.MSELoss() optimizer = optim.Adam(net.parameters(), lr=0.001) # 训练模型 for epoch in range(10): # 迭代 10 次 running_loss = 0.0 for i, data in enumerate(trainloader, 0): inputs, _ = data # 添加椒盐噪声 noise = torch.rand_like(inputs) < 0.02 inputs[noise] = torch.rand_like(inputs[noise]) optimizer.zero_grad() outputs = net(inputs) loss = criterion(outputs, inputs) loss.backward() optimizer.step() running_loss += loss.item() if i % 200 == 199: print('[%d, %5d] loss: %.3f' % (epoch + 1, i + 1, running_loss / 200)) running_loss = 0.0 print('Training finished') # 加载测试集 testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform) testloader = torch.utils.data.DataLoader(testset, batch_size=4, shuffle=False, num_workers=2) dataiter = iter(testloader) images, _ = dataiter.next() # 预测去噪后的图像 outputs = net(images) noisy_images = images + 0.02 * torch.randn_like(images) denoised_images = outputs.detach() + noisy_images # 显示原始图像、噪声图像和去噪后的图像 images = torchvision.utils.make_grid(images) noisy_images = torchvision.utils.make_grid(noisy_images) denoised_images = torchvision.utils.make_grid(denoised_images) images = images / 2 + 0.5 # 反归一化 noisy_images = noisy_images / 2 + 0.5 denoised_images = denoised_images / 2 + 0.5 plt.imshow(np.transpose(images, (1, 2, 0))) plt.title('Original Images') plt.show() plt.imshow(np.transpose(noisy_images, (1, 2, 0))) plt.title('Noisy Images') plt.show() plt.imshow(np.transpose(denoised_images, (1, 2, 0))) plt.title('Denoised Images') plt.show() ``` 这个代码使用 CIFAR10 数据集进行训练和测试,并添加了椒盐噪声。它定义了一个卷积神经网络模型,使用均方差损失函数进行训练,优化器使用 Adam。训练完成后,它使用测试集进行图像去噪,并显示原始图像、噪声图像和去噪后的图像。 ### 回答3: 使用PyTorch基于卷积神经网络进行图像去噪的步骤如下所示: 1. 导入所需的库和模块: import torch import torchvision import torchvision.transforms as transforms import torch.nn as nn import torch.optim as optim import torch.nn.functional as F from torch.utils.data import DataLoader 2. 定义CNN模型: class CNN(nn.Module): def __init__(self): super(CNN, self).__init__() self.conv1 = nn.Conv2d(3, 64, kernel_size=3, padding=1) self.conv2 = nn.Conv2d(64, 64, kernel_size=3, padding=1) self.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=1) self.conv4 = nn.Conv2d(128, 128, kernel_size=3, padding=1) self.fc1 = nn.Linear(128 * 8 * 8, 512) self.fc2 = nn.Linear(512, 10) def forward(self, x): x = F.relu(self.conv1(x)) x = F.relu(self.conv2(x)) x = F.max_pool2d(x, kernel_size=2, stride=2) x = F.relu(self.conv3(x)) x = F.relu(self.conv4(x)) x = F.max_pool2d(x, kernel_size=2, stride=2) x = x.view(x.size(0), -1) x = F.relu(self.fc1(x)) x = self.fc2(x) return x 3. 定义训练和测试函数: def train(model, trainloader, criterion, optimizer): model.train() for inputs, labels in trainloader: optimizer.zero_grad() outputs = model(inputs) loss = criterion(outputs, labels) loss.backward() optimizer.step() def test(model, testloader): model.eval() correct = 0 total = 0 with torch.no_grad(): for inputs, labels in testloader: outputs = model(inputs) _, predicted = torch.max(outputs.data, 1) total += labels.size(0) correct += (predicted == labels).sum().item() accuracy = 100 * correct / total return accuracy 4. 加载CIFAR10数据集和添加噪声: transform = transforms.Compose( [transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]) trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform) trainloader = DataLoader(trainset, batch_size=128, shuffle=True, num_workers=2) testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform) testloader = DataLoader(testset, batch_size=128, shuffle=False, num_workers=2) def add_salt_pepper_noise(img, noise_level): img = img.clone() salt = torch.rand(img.size()) < noise_level / 2 pepper = torch.rand(img.size()) < noise_level / 2 img[salt] = 1.0 img[pepper] = 0.0 return img noise_level = 0.1 for i in range(len(trainset)): img, label = trainset[i] img = add_salt_pepper_noise(img, noise_level) trainset[i] = img, label 5. 实例化模型和优化器: model = CNN() criterion = nn.CrossEntropyLoss() optimizer = optim.SGD(model.parameters(), lr=0.01) num_epochs = 10 6. 执行训练和测试: for epoch in range(num_epochs): train(model, trainloader, criterion, optimizer) acc = test(model, testloader) print(f"Epoch [{epoch+1}/{num_epochs}], Accuracy: {acc}%") 7. 展示图片: import matplotlib.pyplot as plt import numpy as np def imshow(img): img = img / 2 + 0.5 # 从[-1,1]范围转换为[0,1] npimg = img.numpy() plt.imshow(np.transpose(npimg, (1, 2, 0))) plt.show() # 展示噪声图像 dataiter = iter(trainloader) images, labels = dataiter.next() imshow(torchvision.utils.make_grid(images)) # 展示去噪后的图像 outputs = model(images) _, predicted = torch.max(outputs, 1) imshow(torchvision.utils.make_grid(images[predicted == labels])) 通过以上步骤,我们使用PyTorch基于卷积神经网络CIFAR10数据集中的图像进行了去噪操作,并展示了噪声图像和去噪后的图像。

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论 2
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

SimpleZihao

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值