用pytorch写一个十分类的,识别叶片的模型

以下是用PyTorch实现的十分类识别叶片的模型的示例代码:

 
 

pythonCopy code

import torch import torch.nn as nn import torch.optim as optim import torchvision.transforms as transforms from torch.utils.data import DataLoader from torchvision.datasets import ImageFolder # 设置随机种子 torch.manual_seed(123) # 设置超参数 batch_size = 32 learning_rate = 0.001 num_epochs = 10 # 数据预处理 transform = transforms.Compose([ transforms.Resize((224, 224)), # 调整图像大小 transforms.ToTensor(), # 转换为张量 transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) # 归一化 ]) # 加载数据集 train_dataset = ImageFolder(root='train', transform=transform) train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True) test_dataset = ImageFolder(root='test', transform=transform) test_loader = DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False) # 定义模型 class LeafNet(nn.Module): def __init__(self): super(LeafNet, self).__init__() self.conv1 = nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1) self.relu1 = nn.ReLU() self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2) self.conv2 = nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1) self.relu2 = nn.ReLU() self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2) self.conv3 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1) self.relu3 = nn.ReLU() self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2) self.fc1 = nn.Linear(64 * 28 * 28, 512) self.relu4 = nn.ReLU() self.fc2 = nn.Linear(512, 10) def forward(self, x): x = self.conv1(x) x = self.relu1(x) x = self.pool1(x) x = self.conv2(x) x = self.relu2(x) x = self.pool2(x) x = self.conv3(x) x = self.relu3(x) x = self.pool3(x) x = x.view(-1, 64 * 28 * 28) x = self.fc1(x) x = self.relu4(x) x = self.fc2(x) return x # 定义损失函数和优化器 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') net = LeafNet().to(device) criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(net.parameters(), lr=learning_rate) # 训练模型 for epoch in range(num_epochs): running_loss = 0.0 for i, data in enumerate(train_loader, 0): inputs, labels = data inputs, labels = inputs.to(device), labels.to(device) optimizer.zero_grad() outputs = net(inputs) loss = criterion(outputs, labels) loss.backward() optimizer.step() running_loss += loss.item() if i % 100 == 99: print('Epoch [%d/%d], Iteration [%d/%d], Loss: %.4f' % (epoch + 1, num_epochs, i + 1, len(train_loader), running_loss / 100)) running_loss = 0.0 # 测试模型 correct = 0 total = 0 with torch.no_grad(): for data in test_loader: images, labels = data images, labels = images.to(device), labels.to(device) 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 test images: %d %%' % (100 * correct / total))

其中,训练集和测试集都需要按照以下方式组织:

 
 

plaintextCopy code

dataset/ train/ class1/ 1.jpg 2.jpg ... class2/ 1.jpg 2.jpg ... ... test/ class1/ 1.jpg 2.jpg ... class2/ 1.jpg 2.jpg ... ...

其中,class1class2等表示不同的叶片类别,1.jpg2.jpg等表示对应类别的不同样本。在代码中,ImageFolder可以自动将不同文件夹中的图像分为不同的类别,并为每个样本分配一个标签。

猜你喜欢

转载自blog.csdn.net/ihateright/article/details/131134873