PyTorch训练Celeba | Bohrium-玻尔科研空间站

空间站广场

论文

Notebooks

比赛

课程

Apps

我的主页

我的Notebooks

我的论文库

我的足迹

我的工作空间

任务

节点

文件

数据集

镜像

项目

数据库

公开

PyTorch训练Celeba

PyTorch

Celeba

PyTorchCeleba

更新于 2024-10-04

推荐镜像 :pytorch:v1

推荐机型 :c3_m4_1 * NVIDIA T4

数据集

celeba(v1)

CelebFaces属性数据集（CelebA）是一个大规模的人脸属性数据集，包含超过20万张名人图像，每张图像都有40个属性标注。该数据集中的图像涵盖了大范围的姿态变化和复杂的背景。CelebA具有高度的多样性、大量的数据和丰富的标注信息，包括：

10,177个身份，
202,599张人脸图像，
每张图像有5个标志点位置和40个二进制属性标注。

该数据集可用于以下计算机视觉任务的训练和测试集：人脸属性识别、人脸识别、人脸检测、标志点（或面部部位）定位以及人脸编辑与合成。

代码

文本

1. 安装必要的依赖

如果还没有安装PyTorch和Torchvision，可以通过pip安装：

代码

文本

[1]

%pip install torch torchvision

代码

文本

2. 加载CelebA数据集

我们可以通过torchvision.datasets模块轻松获取它。

测试google drive下载困难，已经下好放到数据集里，从数据集复制到/root下并解压

代码

文本

[1]

!cp /bohr/celeba-p5yz/v1/celeba.zip /root

代码

文本

[4]

!cd /root && unzip -q celeba.zip && echo -ne 'Extracting... \r'

代码

文本

[5]

import torch

import torch.nn as nn

import torch.optim as optim

import torchvision

import torchvision.transforms as transforms

from torch.utils.data import DataLoader

# 图像的预处理：缩放到 64x64 大小，并归一化

# 这个变换管道通常用于图像分类任务中的数据预处理步骤。通过调整图像大小、转换为张量和标准化，可以确保输入数据的一致性和模型训练的稳定性。

transform = transforms.Compose([

transforms.Resize((64, 64)),

#将图像从 PIL 图像或 numpy 数组转换为 PyTorch 张量，并且会自动将图像的像素值从 [0, 255] 范围缩放到 [0, 1] 范围。

transforms.ToTensor(),

#使用给定的均值和标准差对图像进行标准化。这里的均值和标准差分别是 [0.5, 0.5, 0.5]，表示对每个通道（RGB）进行标准化。标准化公式为： [ text{output} = frac{text{input} - text{mean}}{text{std}} ] 由于输入的像素值已经被 transforms.ToTensor 缩放到 [0, 1] 范围，标准化后的像素值将被调整到 [-1, 1] 范围。

transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])

])

# 加载 CelebA 数据集,如果google drive可以用，download=True，可以直接下载，也可以自己下载后解压到/root/celeba目录下

root = '/root'

train_dataset = torchvision.datasets.CelebA(root=root, split='train', download=False, transform=transform)

test_dataset = torchvision.datasets.CelebA(root=root, split='test', download=False, transform=transform)

# 创建数据加载器，改大batch_size，在GPU T4下显存使用率没增加？

batch_size = 128

train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)

test_loader = DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False)

# 检查数据加载是否成功

print(f"Number of training samples: {len(train_dataset)}")

print(f"Number of testing samples: {len(test_dataset)}")

Number of training samples: 162770
Number of testing samples: 19962

代码

文本

3. 定义CNN模型

接下来，我们定义一个简单的卷积神经网络模型，结构可以根据任务的复杂性调整。这里的例子是一个基础的CNN。

3.1 初始化

卷积层：

self.conv1：第一个卷积层，输入通道数为 3（RGB 图像），输出通道数为 64，卷积核大小为 3x3，步幅为 1，填充为 1。
self.conv2：第二个卷积层，输入通道数为 64，输出通道数为 128，卷积核大小为 3x3，步幅为 1，填充为 1。
self.conv3：第三个卷积层，输入通道数为 128，输出通道数为 256，卷积核大小为 3x3，步幅为 1，填充为 1。

全连接层：

self.fc1：第一个全连接层，输入大小为 25688，输出大小为 512。
self.fc2：第二个全连接层，输入大小为 512，输出大小为 40（CelebA 数据集有 40 个属性标签）。

池化层和激活函数：

self.pool：最大池化层，池化核大小为 2x2，步幅为 2。
self.relu：ReLU 激活函数。

代码

文本

[6]

class SimpleCNN(nn.Module):

def __init__(self):

super(SimpleCNN, self).__init__()

self.conv1 = nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, stride=1, padding=1)

self.conv2 = nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, stride=1, padding=1)

self.conv3 = nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, stride=1, padding=1)

self.fc1 = nn.Linear(256*8*8, 512)

self.fc2 = nn.Linear(512, 40) # CelebA 有 40 个属性标签

self.pool = nn.MaxPool2d(kernel_size=2, stride=2)

self.relu = nn.ReLU()

def forward(self, x):

x = self.relu(self.conv1(x))

x = self.pool(x)

x = self.relu(self.conv2(x))

x = self.pool(x)

x = self.relu(self.conv3(x))

x = self.pool(x)

x = x.view(x.size(0), -1) # 展平操作

x = self.relu(self.fc1(x))

x = self.fc2(x)

return x

# 实例化模型

model = SimpleCNN()

代码

文本

4. 定义损失函数和优化器

对于多标签分类任务，我们使用 BCEWithLogitsLoss 损失函数。优化器可以选择 Adam 或 SGD。

代码

文本

[7]

criterion = nn.BCEWithLogitsLoss() # 二分类交叉熵损失

optimizer = optim.Adam(model.parameters(), lr=0.001)

代码

文本

5. 训练模型

下面是训练模型的代码，每个 epoch 后对模型进行验证。

5.1. train 函数

train 函数用于训练模型。

参数解释

model：要训练的神经网络模型。
train_loader：训练数据的 DataLoader 对象。
criterion：损失函数，用于计算预测值和真实值之间的误差。
optimizer：优化器，用于更新模型的参数。
num_epochs：训练的轮数。
device：计算设备（如 CPU 或 GPU）。

5.2. evaluate 函数

evaluate 函数用于评估模型的性能。

参数解释

model：要评估的神经网络模型。
test_loader：测试数据的 DataLoader 对象。
criterion：损失函数，用于计算预测值和真实值之间的误差。
device：计算设备（如 CPU 或 GPU）。

代码

文本

[8]

def train(model, train_loader, criterion, optimizer, num_epochs, device):

model.train() #设置模型为训练模式

for epoch in range(num_epochs):

for images, labels in train_loader:

# 将数据移动到选定的设备

images, labels = images.to(device), labels.to(device).float()

# 前向传播

outputs = model(images)

loss = criterion(outputs, labels)

# 反向传播和优化

optimizer.zero_grad()#清零梯度

loss.backward()#计算梯度

optimizer.step()#更新参数

print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}')

def evaluate(model, test_loader, criterion, device):

model.eval() # 设置模型为评估模式

with torch.no_grad(): # 禁用梯度计算

total_loss = 0

correct = 0

total = 0

for images, labels in test_loader:

# 将数据移动到选定的设备

images, labels = images.to(device), labels.to(device).float()

# 前向传播

outputs = model(images)

# 计算损失，并累加到 total_loss

loss = criterion(outputs, labels)

total_loss += loss.item()

# 计算准确率

predicted = (outputs > 0.5).float()

total += labels.size(0) * labels.size(1) # 总标签数

correct += (predicted == labels).sum().item()

# 打印平均损失和准确率

print(f'Average loss: {total_loss / len(test_loader):.4f}, Accuracy: {100 * correct / total:.2f}%')

代码

文本

开始训练

通过调用train函数进行训练，完成后用evaluate函数进行验证。

代码

文本

[9]

# 检查 CUDA 是否可用，并选择设备

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 将模型移动到选定的设备

model = model.to(device)

# 训练模型

train(model, train_loader, criterion, optimizer, num_epochs=10, device=device)

# 验证模型

evaluate(model, test_loader, criterion, device=device)

Epoch [1/10], Loss: 0.2363
Epoch [2/10], Loss: 0.2335
Epoch [3/10], Loss: 0.2055
Epoch [4/10], Loss: 0.2028
Epoch [5/10], Loss: 0.2045
Epoch [6/10], Loss: 0.1880
Epoch [7/10], Loss: 0.1708
Epoch [8/10], Loss: 0.1693
Epoch [9/10], Loss: 0.1608
Epoch [10/10], Loss: 0.1483
Average loss: 0.2647, Accuracy: 89.48%

代码

文本

保存和加载模型

训练完成后，可以保存模型以便之后使用。

代码

文本

[10]

# 保存模型

torch.save(model.state_dict(), 'celeba_cnn.pth')

# 加载模型

model.load_state_dict(torch.load('celeba_cnn.pth'))

<All keys matched successfully>

代码

文本

PyTorch

Celeba

PyTorchCeleba

点个赞吧