CelebFaces属性数据集（CelebA）是一个大规模的人脸属性数据集，包含超过20万张名人图像，每张图像都有40个属性标注。该数据集中的图像涵盖了大范围的姿态变化和复杂的背景。CelebA具有高度的多样性、大量的数据和丰富的标注信息，包括：
- 10,177个身份，
- 202,599张人脸图像，
- 每张图像有5个标志点位置和40个二进制属性标注。
该数据集可用于以下计算机视觉任务的训练和测试集：人脸属性识别、人脸识别、人脸检测、标志点（或面部部位）定位以及人脸编辑与合成。
1. 安装必要的依赖
如果还没有安装PyTorch和Torchvision，可以通过pip安装：
%pip install torch torchvision
2. 加载CelebA数据集
我们可以通过torchvision.datasets模块轻松获取它。
如果没有数据集且无法通过google drive下载，可以通过这篇notebook关联的数据集下载
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)}")
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 激活函数。
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。
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）。
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函数进行验证。
# 检查 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)
- 保存和加载模型
训练完成后，可以保存模型以便之后使用。
# 保存模型
torch.save(model.state_dict(), 'celeba_cnn.pth')
# 加载模型
model.load_state_dict(torch.load('celeba_cnn.pth'))
