Deepfake-FFDI-图像赛题 baseline（2）模型训练与验证 部分 @Datawhale AI 夏令营

class AverageMeter(object):
    """Computes and stores the average and current value"""
    def __init__(self, name, fmt=':f'):
        self.name = name
        self.fmt = fmt
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count

    def __str__(self):
        fmtstr = '{name} {val' + self.fmt + '} ({avg' + self.fmt + '})'
        return fmtstr.format(**self.__dict__)

class ProgressMeter(object):
    def __init__(self, num_batches, *meters):
        self.batch_fmtstr = self._get_batch_fmtstr(num_batches)
        self.meters = meters
        self.prefix = ""


    def pr2int(self, batch):
        entries = [self.prefix + self.batch_fmtstr.format(batch)]
        entries += [str(meter) for meter in self.meters]
        print('\t'.join(entries))

    def _get_batch_fmtstr(self, num_batches):
        num_digits = len(str(num_batches // 1))
        fmt = '{:' + str(num_digits) + 'd}'
        return '[' + fmt + '/' + fmt.format(num_batches) + ']'

 这段代码定义了两个类AverageMeter 和 ProgressMeter，它们主要用于在训练或验证神经网络模型时，方便地记录和显示性能指标（例如损失、精度等）的平均值和进度。

AverageMeter 类

AverageMeter 类用于计算和存储一个指标的当前值、累计值、总和和平均值。这个类通常用于记录损失值、精度等性能指标。

方法和属性解释：

1. __init__(self, name, fmt=':f')：

   • 初始化方法，接收指标的名称 name 和格式字符串 fmt（默认是浮点数格式 ':f'）。
   • 调用 reset() 方法重置所有计数器。

2. reset(self)：

   • 重置所有内部状态，包括当前值 (val)、平均值 (avg)、总和 (sum) 和计数 (count)。

3. update(self, val, n=1)：

   • 更新当前值 val，并根据给定的数量 n 更新总和和计数器，然后计算新的平均值。

4. __str__(self)：

   • 返回格式化的字符串，显示当前值和平均值。

ProgressMeter 类

ProgressMeter 类用于显示训练过程中的进度和性能指标。它会在每个批次完成时打印当前的批次编号和各个指标的当前值及平均值。

方法和属性解释：

1. __init__(self, num_batches, *meters)：

   • 初始化方法，接收总的批次数 num_batches 和若干个 AverageMeter 实例。
   • 使用 _get_batch_fmtstr(num_batches) 方法生成批次格式字符串，并将其存储在 batch_fmtstr 属性中。
   • 将 meters 参数保存到 meters 属性中，并初始化 prefix 属性为空字符串。

2. pr2int(self, batch)：

   • 打印当前批次的进度和所有 AverageMeter 实例的状态。
   • 使用批次格式字符串格式化当前批次编号，并将各个 AverageMeter 实例的字符串表示拼接起来打印。

3. _get_batch_fmtstr(self, num_batches)：

   • 根据总的批次数 num_batches 生成格式字符串。
   • 计算批次数的位数，生成形如 [1/100] 的格式字符串，用于打印进度。

 

validate 函数

validate 函数用于在验证集上评估模型的性能。它会计算模型的损失和准确率，并返回平均准确率。

1.初始化性能指标

batch_time = AverageMeter('Time', ':6.3f')
losses = AverageMeter('Loss', ':.4e')
top1 = AverageMeter('Acc@1', ':6.2f')
progress = ProgressMeter(len(val_loader), batch_time, losses, top1)

• 创建三个 AverageMeter 实例来记录时间、损失和准确率。
• 创建一个 ProgressMeter 实例来显示进度。

2.切换模型到评估模式：

model.eval()

 3.禁用梯度计算：

with torch.no_grad():

4. 遍历验证集：

for i, (input, target) in tqdm_notebook(enumerate(val_loader), total=len(val_loader)):
    input = input.cuda()
    target = target.cuda()
    output = model(input)
    loss = criterion(output, target)
    acc = (output.argmax(1).view(-1) == target.float().view(-1)).float().mean() * 100
    losses.update(loss.item(), input.size(0))
    top1.update(acc, input.size(0))
    batch_time.update(time.time() - end)
    end = time.time()

5.最终打印平均准确率

print(' * Acc@1 {top1.avg:.3f}'.format(top1=top1))
return top1

 

predict 函数

predict 函数用于在测试集上进行预测，并支持 TTA（Test Time Augmentation）。

1.切换模型到评估模式

model.eval()

 2.初始化一个存储预测结果的变量

test_pred_tta = None

 3.执行TTA循环

for _ in range(tta):
    test_pred = []
    with torch.no_grad():
        for i, (input, target) in tqdm_notebook(enumerate(test_loader), total=len(test_loader)):
            input = input.cuda()
            output = model(input)
            output = F.softmax(output, dim=1)
            output = output.data.cpu().numpy()
            test_pred.append(output)
    test_pred = np.vstack(test_pred)
    if test_pred_tta is None:
        test_pred_tta = test_pred
    else:
        test_pred_tta += test_pred

 4.返回累加的预测结果

return test_pred_tta

 

train 函数

train 函数用于在训练集上训练模型。它会计算模型的损失和准确率，并更新模型的参数。

与上面类似，不在赘述。代码如下：

#初始化性能指标
batch_time = AverageMeter('Time', ':6.3f')
losses = AverageMeter('Loss', ':.4e')
top1 = AverageMeter('Acc@1', ':6.2f')
progress = ProgressMeter(len(train_loader), batch_time, losses, top1)

#切换模型到训练模式
model.train()

#遍历数据集
for i, (input, target) in enumerate(train_loader):
    input = input.cuda(non_blocking=True)
    target = target.cuda(non_blocking=True)
    output = model(input)
    loss = criterion(output, target)
    losses.update(loss.item(), input.size(0))
    acc = (output.argmax(1).view(-1) == target.float().view(-1)).float().mean() * 100
    top1.update(acc, input.size(0))
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    batch_time.update(time.time() - end)
    end = time.time()
    if i % 100 == 0:
        progress.pr2int(i)

下面的 这段代码定义了一个自定义的 Dataset 类，名为 FFDIDataset，用于处理图像数据和相应的标签。这个类继承了 torch.utils.data.Dataset，它是 PyTorch 中用于加载数据的基类。

class FFDIDataset(Dataset):
    def __init__(self, img_path, img_label, transform=None):
        self.img_path = img_path
        self.img_label = img_label
        
        if transform is not None:
            self.transform = transform
        else:
            self.transform = None
    
    def __getitem__(self, index):
        img = Image.open(self.img_path[index]).convert('RGB')
        
        if self.transform is not None:
            img = self.transform(img)
        
        return img, torch.from_numpy(np.array(self.img_label[index]))
    
    def __len__(self):
        return len(self.img_path)

1. __init__(self, img_path, img_label, transform=None)：

   • 初始化方法，用于设置数据集的图像路径、标签和数据转换。
   • 参数解释：
     • img_path：一个列表，包含所有图像文件的路径。
     • img_label：一个列表，包含所有图像对应的标签。
     • transform：一个可选参数，用于图像的预处理和数据增强。如果没有提供，默认为 None。
   • 在初始化方法中，将 img_path 和 img_label 分别赋值给类的属性 self.img_path 和 self.img_label。同时，检查是否提供了 transform，如果提供了就赋值给 self.transform，否则设置为 None。

2. __getitem__(self, index)：

   • 用于根据索引获取数据集中的一个样本。
   • 参数解释：
     • index：数据集中样本的索引。
   • 方法步骤：
     1. 使用 Image.open 打开索引 index 处的图像文件，并将其转换为 RGB 格式。
     2. 如果提供了 transform，对图像进行预处理或数据增强。
     3. 返回处理后的图像和对应的标签，标签使用 torch.from_numpy 转换为 PyTorch 张量。

3. __len__(self)：

   • 返回数据集中样本的数量，即图像路径列表的长度。