Python基于PyTorch实现卷积神经网络回归模型
·
下面是一个基于PyTorch实现的卷积神经网络回归模型的示例代码:
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision.datasets import CIFAR10
from torchvision.transforms import ToTensor, Normalize
from tqdm import tqdm
# 定义卷积神经网络模型
class CNNRegressor(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
self.bn1 = nn.BatchNorm2d(32)
self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
self.bn2 = nn.BatchNorm2d(64)
self.conv3 = nn.Conv2d(64, 128, 3, padding=1)
self.bn3 = nn.BatchNorm2d(128)
self.fc1 = nn.Linear(128*4*4, 512)
self.fc2 = nn.Linear(512, 1)
def forward(self, x):
x = nn.functional.relu(self.bn1(self.conv1(x)))
x = nn.functional.max_pool2d(x, 2)
x = nn.functional.relu(self.bn2(self.conv2(x)))
x = nn.functional.max_pool2d(x, 2)
x = nn.functional.relu(self.bn3(self.conv3(x)))
x = nn.functional.max_pool2d(x, 2)
x = x.view(-1, 128*4*4)
x = nn.functional.relu(self.fc1(x))
x = self.fc2(x)
return x
# 加载数据集
train_dataset = CIFAR10(root='./data', train=True, download=True, transform=Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]))
train_dataloader = DataLoader(train_dataset, batch_size=128, shuffle=True)
test_dataset = CIFAR10(root='./data', train=False, download=True, transform=Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]))
test_dataloader = DataLoader(test_dataset, batch_size=128, shuffle=False)
# 定义模型、损失函数、优化器
model = CNNRegressor()
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
# 训练模型
num_epochs = 10
for epoch in range(num_epochs):
model.train()
train_loss = 0.0
for inputs, targets in tqdm(train_dataloader, desc=f"Epoch {epoch+1}/{num_epochs}", unit="batch"):
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, targets.float().unsqueeze(1))
loss.backward()
optimizer.step()
train_loss += loss.item() * inputs.size(0)
train_loss /= len(train_dataset)
model.eval()
test_loss = 0.0
with torch.no_grad():
for inputs, targets in tqdm(test_dataloader, desc="Testing", unit="batch"):
outputs = model(inputs)
loss = criterion(outputs, targets.float().unsqueeze(1))
test_loss += loss.item() * inputs.size(0)
test_loss /= len(test_dataset)
print(f"Train Loss: {train_loss:.4f}, Test Loss: {test_loss:.4f}")
这个示例代码使用了CIFAR10数据集,卷积神经网络模型包括3个卷积层和2个全连接层,优化器使用Adam,训练10个epoch。在训练期间,使用了PyTorch的DataLoader和tqdm库来进行数据迭代和进度条显示,同时计算训练集和测试集的损失。
更多推荐
所有评论(0)