写一个pythorch框架的,用AlexNet进行128*128图像的4分类的程序
时间: 2025-01-29 18:20:24 浏览: 38
以下是一个使用PyTorch框架和AlexNet模型进行128x128图像4分类的示例程序:
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms, models
# 定义数据预处理
transform = transforms.Compose([
transforms.Resize((128, 128)), # 将图像大小调整为128x128
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], # 使用ImageNet的均值和标准差
[0.229, 0.224, 0.225])
])
# 加载数据集
train_dataset = datasets.ImageFolder(root='path_to_train_data', transform=transform)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=32, shuffle=True)
val_dataset = datasets.ImageFolder(root='path_to_val_data', transform=transform)
val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=32, shuffle=False)
# 加载预训练的AlexNet模型
model = models.alexnet(pretrained=True)
# 修改最后的全连接层以适应4分类任务
num_ftrs = model.classifier[6].in_features
model.classifier[6] = nn.Linear(num_ftrs, 4)
# 移动模型到GPU(如果可用)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = model.to(device)
# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
# 训练模型
num_epochs = 10
for epoch in range(num_epochs):
model.train()
running_loss = 0.0
for inputs, labels in train_loader:
inputs = inputs.to(device)
labels = labels.to(device)
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item() * inputs.size(0)
epoch_loss = running_loss / len(train_dataset)
# 验证模型
model.eval()
correct = 0
total = 0
with torch.no_grad():
for inputs, labels in val_loader:
inputs = inputs.to(device)
labels = labels.to(device)
outputs = model(inputs)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
accuracy = correct / total
print(f'Epoch {epoch+1}/{num_epochs}, Loss: {epoch_loss:.4f}, Validation Accuracy: {accuracy:.4f}')
print('Training complete')
这个程序包含了以下主要部分:
- 数据预处理:调整图像大小,转换为张量,并进行标准化。
- 数据加载:使用ImageFolder和DataLoader加载训练和验证数据。
- 模型定义:加载预训练的AlexNet模型,并修改最后的全连接层以适应4分类任务。
- 训练过程:定义损失函数和优化器,并进行多轮训练。
- 验证过程:在每个epoch后评估模型在验证集上的性能。
请注意,你需要将'path_to_train_data'和'path_to_val_data'替换为你的实际数据路径。