pytorch深度学习猫狗识别代码

以下是一个基于PyTorch的猫狗识别代码示例：

首先，我们需要导入必要的库和设置随机数种子：

import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
import numpy as np
import random

# 设置随机数种子
random.seed(0)
np.random.seed(0)
torch.manual_seed(0)
torch.cuda.manual_seed(0)

接着，我们需要加载数据集并进行数据增强：

# 定义数据增强
transform_train = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

transform_test = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# 加载数据集
train_dataset = torchvision.datasets.ImageFolder(root='./data/train', transform=transform_train)
test_dataset = torchvision.datasets.ImageFolder(root='./data/test', transform=transform_test)

# 定义数据加载器
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=2)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=32, shuffle=False, num_workers=2)

然后，我们需要定义一个卷积神经网络模型：

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm2d(32)
        self.relu1 = nn.ReLU()
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm2d(64)
        self.relu2 = nn.ReLU()
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
        self.bn3 = nn.BatchNorm2d(128)
        self.relu3 = nn.ReLU()
        self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = nn.Linear(128 * 28 * 28, 512)
        self.bn4 = nn.BatchNorm1d(512)
        self.relu4 = nn.ReLU()
        self.fc2 = nn.Linear(512, 2)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu1(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.relu2(x)
        x = self.pool2(x)
        x = self.conv3(x)
        x = self.bn3(x)
        x = self.relu3(x)
        x = self.pool3(x)
        x = x.view(-1, 128 * 28 * 28)
        x = self.fc1(x)
        x = self.bn4(x)
        x = self.relu4(x)
        x = self.fc2(x)
        return x

接着，我们需要定义损失函数和优化器：

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = CNN().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

接下来，我们可以开始训练模型了：

num_epochs = 10

for epoch in range(num_epochs):
    # 训练模型
    model.train()
    train_loss = 0.0
    train_acc = 0.0
    for images, labels in train_loader:
        images = images.to(device)
        labels = labels.to(device)
        outputs = model(images)
        loss = criterion(outputs, labels)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        train_loss += loss.item() * images.size(0)
        _, preds = torch.max(outputs, 1)
        train_acc += torch.sum(preds == labels.data)
    train_loss = train_loss / len(train_dataset)
    train_acc = train_acc / len(train_dataset)

    # 测试模型
    model.eval()
    test_loss = 0.0
    test_acc = 0.0
    for images, labels in test_loader:
        images = images.to(device)
        labels = labels.to(device)
        outputs = model(images)
        loss = criterion(outputs, labels)
        test_loss += loss.item() * images.size(0)
        _, preds = torch.max(outputs, 1)
        test_acc += torch.sum(preds == labels.data)
    test_loss = test_loss / len(test_dataset)
    test_acc = test_acc / len(test_dataset)

    print('Epoch [{}/{}], Train Loss: {:.4f}, Train Acc: {:.4f}, Test Loss: {:.4f}, Test Acc: {:.4f}'.format(
        epoch+1, num_epochs, train_loss, train_acc, test_loss, test_acc))

最后，我们可以使用训练好的模型进行预测：

# 定义类别名称
classes = ['cat', 'dog']

# 读取一张测试图片
image = Image.open('./data/test/cat/cat.1000.jpg')

# 对图片进行预处理
image_tensor = transform_test(image).float()
image_tensor = image_tensor.unsqueeze_(0)

# 使用训练好的模型进行预测
model.eval()
with torch.no_grad():
    outputs = model(image_tensor)
    _, preds = torch.max(outputs, 1)

# 输出预测结果
print(classes[preds.item()])