用pytorch写猫狗识别代码
时间: 2023-12-23 22:04:18 浏览: 85
好的,以下是一个简单的猫狗识别代码,基于PyTorch框架实现:
首先,我们需要导入相关的库:
```python
import os
import torch
import torchvision
import torch.nn as nn
import torch.optim as optim
import numpy as np
from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder
from torchvision.transforms import transforms
```
然后,我们需要定义一些超参数:
```python
batch_size = 32
num_epochs = 10
learning_rate = 0.001
```
接着,我们需要定义数据预处理的方式:
```python
transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
```
这里我们使用了 `Resize` 来将图片的大小调整为 (224, 224),使用 `ToTensor` 将图片转换为张量,并使用 `Normalize` 对张量进行标准化处理。
然后,我们需要加载数据集:
```python
train_dataset = ImageFolder('./data/train', transform=transform)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_dataset = ImageFolder('./data/test', transform=transform)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
```
这里我们使用了 `ImageFolder` 来加载数据集。`train_dataset` 和 `test_dataset` 分别表示训练集和测试集。`train_loader` 和 `test_loader` 分别表示训练集和测试集的数据加载器。
接着,我们需要定义一个卷积神经网络模型:
```python
class CatDogClassifier(nn.Module):
def __init__(self):
super(CatDogClassifier, self).__init__()
self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1)
self.conv2 = nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1)
self.conv3 = nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1)
self.conv4 = nn.Conv2d(256, 512, kernel_size=3, stride=1, padding=1)
self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
self.fc1 = nn.Linear(512 * 28 * 28, 1024)
self.fc2 = nn.Linear(1024, 512)
self.fc3 = nn.Linear(512, 2)
self.dropout = nn.Dropout(p=0.5)
def forward(self, x):
x = self.pool(torch.relu(self.conv1(x)))
x = self.pool(torch.relu(self.conv2(x)))
x = self.pool(torch.relu(self.conv3(x)))
x = self.pool(torch.relu(self.conv4(x)))
x = x.view(-1, 512 * 28 * 28)
x = self.dropout(torch.relu(self.fc1(x)))
x = self.dropout(torch.relu(self.fc2(x)))
x = self.fc3(x)
return x
```
这里我们定义了一个四层卷积神经网络模型,包括两个全连接层。具体结构如下:
- 卷积层1:输入通道数为3,输出通道数为64,卷积核大小为3x3。
- 卷积层2:输入通道数为64,输出通道数为128,卷积核大小为3x3。
- 卷积层3:输入通道数为128,输出通道数为256,卷积核大小为3x3。
- 卷积层4:输入通道数为256,输出通道数为512,卷积核大小为3x3。
- 最大池化层:池化核大小为2x2,步长为2。
- 全连接层1:输入大小为512x28x28,输出大小为1024。
- 全连接层2:输入大小为1024,输出大小为512。
- 全连接层3:输入大小为512,输出大小为2。
接着,我们需要定义损失函数和优化器:
```python
model = CatDogClassifier()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)
```
这里我们使用了交叉熵损失函数和Adam优化器。
最后,我们可以开始训练模型:
```python
for epoch in range(num_epochs):
for i, (images, labels) in enumerate(train_loader):
optimizer.zero_grad()
outputs = model(images)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
if (i + 1) % 100 == 0:
print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'
.format(epoch + 1, num_epochs, i + 1, len(train_loader), loss.item()))
# 测试模型
model.eval()
with torch.no_grad():
correct = 0
total = 0
for images, labels in test_loader:
outputs = model(images)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
print('Test Accuracy of the model on the test images: {} %'.format(100 * correct / total))
model.train()
```
这里我们使用了两层循环进行训练。在每个epoch中,我们遍历训练集,计算损失,并更新参数。每100个batch打印一次损失。然后在测试集上测试模型的准确率。
完整代码如下:
阅读全文