Fairlearn与PyTorch模型结合使用的技术探索

背景介绍

Fairlearn是一个用于评估和改善机器学习模型公平性的Python工具包。虽然它原生支持scikit-learn风格的模型接口，但许多开发者希望将其与PyTorch深度学习框架结合使用。本文将探讨如何通过skorch库实现这一目标，并分享实际应用中的技术细节。

技术挑战

Fairlearn设计时主要考虑了scikit-learn风格的API，而PyTorch模型通常不直接支持sample_weight参数，这给公平性约束的实现带来了挑战。主要问题包括：

1. PyTorch模型默认不处理样本权重
2. 训练和测试阶段需要不同的权重处理逻辑
3. 损失函数需要特殊处理以支持公平性约束

解决方案实现

自定义PyTorch模型

通过扩展nn.Module类，我们可以创建一个支持样本权重的自定义模型：

class FairlearnMLP(nn.Module):
    def __init__(self, num_classes, num_features, num_layers=4, layer_width=128):
        super().__init__()
        layers = []
        for i in range(num_layers - 1):
            layers.append(nn.Linear(num_features if i == 0 else layer_width, layer_width))
            layers.append(nn.ReLU())
        layers.append(nn.Linear(layer_width, num_classes))
        layers.append(nn.Sigmoid())
        self.layers = nn.Sequential(*layers)
        
        self.train_sample_weight = None
        self.test_sample_weight = None
        self.training_samples = training_samples
        self.testing_samples = testing_samples

    def forward(self, x, *, sample_weight=None):
        if sample_weight is not None:
            if len(sample_weight.values) == self.training_samples:
                self.train_sample_weight = sample_weight.values
            elif len(sample_weight.values) == self.testing_samples:
                self.test_sample_weight = sample_weight.values
        return self.layers(x)

自定义skorch分类器

通过扩展NeuralNetBinaryClassifier，我们可以实现样本权重在损失函数中的正确应用：

class FairClassifier(NeuralNetBinaryClassifier):
    def predict(self, X):
        predictions = self.predict_proba(X)
        return (predictions > self.threshold).astype(int)

    def get_loss(self, y_pred, y_true, X=None, training=False):
        y_true = torch.tensor(y_true, dtype=torch.double, device=self.device)
        
        if hasattr(self.module_, 'train_sample_weight') and self.module_.train_sample_weight is not None:
            sample_weight = self.module_.train_sample_weight
        elif hasattr(self.module_, 'test_sample_weight') and self.module_.test_sample_weight is not None:
            sample_weight = self.module_.test_sample_weight
        else:
            sample_weight = None

        if sample_weight is not None:
            loss = self.criterion(y_pred, y_true)
            loss = loss * torch.tensor(sample_weight, dtype=torch.double)
            return loss.mean()
        return self.criterion_(y_pred, y_true).mean()

公平性约束的实现

在Fairlearn中，EqualizedOdds矩约束是最接近以下数学表达式的实现：

$$\Bigg | \frac{1}{n_0} \sum_{i = 1}^{n_0} l(y_{i, 0}, g_{\theta}(\xi)) - \frac{1}{n_1} \sum_{i = 1}^{n_1} l(y_{i, 1}, g_{\theta}(\xi)) \Bigg | \leq r$$

EqualizedOdds约束确保模型在不同敏感属性组上的真正例率和假正例率差异不超过阈值r。对于二分类问题，这是最合适的公平性约束之一。

实践建议

1. 批量大小选择：建议使用全批量训练以确保样本权重正确应用
2. 数据类型一致性：确保权重张量与模型输出使用相同的数据类型
3. 训练测试分离：明确区分训练和测试阶段的样本权重处理
4. 公平性评估：定期使用Fairlearn的指标评估模型公平性

总结

虽然Fairlearn与PyTorch的直接集成需要额外的工作，但通过skorch和自定义模型类可以实现有效的结合。这种方法为深度学习模型提供了公平性保障，同时保持了PyTorch的灵活性和性能优势。开发者可以根据具体需求调整模型结构和公平性约束，构建既准确又公平的机器学习系统。