简介这套基于PyTorch的生成对抗网络缺失数据填补代码包面向需要处理数据缺失问题的机器学习算法工程师、数据科学从业者及对GAIN模型感兴趣的在校研究生。压缩包共27个文件大小约6.72MB除5个Python源码外还包含10个csv格式的公开数据集、PyCharm工程配置XML、编译缓存pyc及README说明文档目录结构清晰便于直接运行、复现实验或在现有基础上做二次开发。代码覆盖GAIN、SGAIN、WSGAIN-CP、WSGAIN-GP四种主流填补方法配套letter、yeast、credit、breast等十个数据集可系统对比不同生成器与损失函数在缺失数据场景下的表现也方便读者快速理解各类变体的实现差异。目前已有2191人学习下载适合希望快速上手生成对抗网络数据填补实验、深入理解模型训练与评估流程的读者是一份兼顾原理学习与工程实践的完整参考。 做数据分析的人迟早会撞上同一个问题手里的表格总有那么几列缺数据。以前我处理缺失值第一反应就是均值、中位数顶上去复杂一点的用多重插补跑一晚上但效果始终差点意思。后来接触到基于生成对抗网络的缺失数据填补方法GAIN才意识到“填缺失”这件事本质上是在学习数据的联合分布而不是在“猜”一个数。这篇文章我就把GAIN在PyTorch下的完整实现从头到尾拆一遍包括网络怎么搭、损失函数怎么配、训练有哪些坑全部摊开讲代码也是完整版的可以直接拿去改自己的数据。1. 缺失值为什么难填传统方法在复杂数据上失灵的本质1.1 均值填补压缩了方差还破坏了变量关系先聊一个最容易被忽略的问题缺失值填补不是“把空位补齐”这么简单它背后是在做“保持数据分布不变”这件事。均值、中位数、众数这类单值填补本质是把所有缺失样本都指向同一个点这会让填补后的变量方差明显变小变量之间的协方差结构也被扭曲。举个直观例子假设身高和体重高度相关你用平均身高填补缺失的部分那这些被填出来的“样本”会全部落在均值竖线上散点图里莫名其妙多出一条直线后续训练聚类、回归、分类模型都会因为这条“假直线”而偏移。1.2 回归填补与多重插补的线性假设天花板稍微讲究一点的做法是回归填补和多重插补MICE。它们把缺失变量当成因变量用其他变量做线性回归来预测缺失值。听起来合理但实际跑过就知道它们的假设是变量关系至少可被线性或低阶非线性描述。真实场景里的用户行为数据、传感器数据、医疗检验数据变量之间的关系往往高度非线性且存在交互效应。你在这种数据上用MICE每插补一次就在线性假设下把误差往下游传一次链式传播到最后误差已经不是“稍微偏一点”而是“系统性偏移”。1.3 GAIN真正在学习的东西P(X | X_observed)GAIN之所以能跳出这些框架是因为它不假设任何显式分布也不假设线性关系。它让生成器去学习“给定已经观察到的部分缺失部分最合理的取值是什么”这个条件分布。生成器见过大量完美的完整样本在对抗训练中被逼着输出让判别器无法分辨真伪的填补值这其实是让模型牢牢抓住了数据的联合分布。换句话说传统方法是在“猜单点”GAIN是在“学习整张分布的形态”这也是我后来在各种表格数据上测试GAIN的RMSE和下游任务表现普遍优于MICE的根本原因。2. 从普通GAN到GAIN生成器、判别器与Hint提示机制的逐层拆解2.1 数据与掩码最不该出错的一步GAIN的输入除了数据矩阵 X还有一个同样重要的矩阵 M叫掩码矩阵。M 中 1 表示该位置被观察到0 表示缺失。这个掩码贯穿网络前向传播和损失计算的每一步很多初版实现跑不出效果问题往往就出在掩码没有同步处理好。我一般在构造输入时把原始数据 X 和随机噪声 Z 按掩码拼接X_tilde M \odot X (1 - M) \odot Z也就是说缺失位置先用随机噪声占位观察位置保持原值。这个 X_tilde 再和 M 拼在一起作为生成器的输入。这样生成器从一开始就知道哪些位置需要修复哪些位置的数据是可信的。这个拼接不是可有可无的设计如果只把 X_tilde 塞给生成器网络就要自己从数值里推断掩码推断错了整个填补就偏了。2.2 生成器一个带掩码条件的修复网络生成器的结构本质上是一个多层全连接网络不需要花里胡哨的卷积或注意力。输入维度是 2dX_tilde 的 d 维加 M 的 d 维输出维度是 d每个特征位对应一个修复后的值。关键在输出层如果数据做了标准化输出层可以用线性激活或 Tanh如果数据是 0-1 区间用 Sigmoid 再缩放到原始范围。实际操作里我给输出层接了线性层然后对输出做了 Clamp把数值限制在训练集特征的最小最大值区间内避免生成器输出极端值拉低评估指标。生成器的输出 X_hat 还不能直接当最终结果要再做一步“合并”X_complete M \odot X (1 - M) \odot X_hat观察位置保留原始值缺失位置用生成值这是 GAIN 的硬性规则。这个合并操作保证模型无论怎么训练都不会去篡改已经观察到的数据这是它与普通自编码器填补的最大区别。2.3 判别器与Hint机制为什么不能让它直接看完整数据GAIN 的判别器输入同样是拼接向量X_complete 拼上 Hint 矩阵 H输出是对掩码矩阵的逐位预测。训练目标是让判别器学会区分“哪些位置是观察到的、哪些是被填补的”而生成器的目标恰恰相反希望判别器在缺失位置出错。这里有一个反直觉的设计Hint 提示矩阵。如果直接让判别器看 X_complete它太容易区分真实值和填补值了因为生成器早期的输出和真实数据差异很大判别器会快速收敛到一个“绝对正确”的状态之后生成器从判别器那里拿不到任何有效梯度训练直接停滞。Hint 机制的做法是B ~ Bernoulli(hint_rate) H M \odot B 0.5 \odot (1 - B)解释成人话对一部分数据点把真实的掩码信息告诉判别器对另一部分数据点给一个 0.5 的模糊值。这样判别器只能部分依赖掩码提示必须学会从数据本身判断缺失模式生成器也不至于被秒杀。2.4 Hint设计失误的现场hint_rate过高会怎样我在一次对比实验里把 hint_rate 调到 0.99判别器几乎完全掌握了真实掩码生成器训练几千轮后依然只会输出数据均值附近的值损失曲线看起来挺平稳但事实就是模型没学到任何条件分布。后来把 hint_rate 降回 0.9训练几十轮后生成器的重建损失明显下降。这说明 Hint 机制不是锦上添花而是 GAIN 能不能有效训练的命门。3. 损失函数设计与训练稳定性调好对抗与重建的平衡3.1 三个损失项各管什么事GAIN 的损失函数有三个来源理解每个来源的作用比照抄公式重要得多。第一是判别器损失它衡量的是判别器预测掩码和真实掩码之间的二分类交叉熵。判别器的目标是把观察位置和填补位置区分开所以它要让这个损失尽可能小。第二是生成器的对抗损失方向相反生成器希望判别器把填补位置也猜成“观察到”也就是让判别器的预测结果趋向全 1 的矩阵。第三是重建损失只对观察位置计算生成器输出和原始值的均方误差它的作用是约束生成器不要为了骗过判别器而随意改写已知信息。用一句话概括对抗损失管“填补得像真的”重建损失管“观察到的别乱动”。两者必须同时存在缺一个模型都会崩。3.2 alpha系数怎么调从1到100我踩过的档位重建损失前面要乘一个系数 alpha用来调节它和对抗损失之间的权重。原论文给的经验值是 alpha 10但不同数据分布差异很大不能死搬。我在 MNIST 上测过 alpha 从 1 到 100 的不同取值观察到的规律是alpha 太小生成器疯狂迎合判别器填出来的数据方差大、形状怪异alpha 太大生成器只顾着把观察位置的重建误差压到最低搞得填补位置全变成均值附近的值跟简单均值填补差不多。我的建议是先从 10 起步观察前 100 个 batch 的重建训练损失如果震荡幅度超过 5%就把 alpha 往上抬一抬如果损失降得很快但数据分布明显偏窄就往下调。3.3 训练崩溃的典型表现与干预手段对抗网络训练不稳定是常态GAIN 也不例外。最典型的崩溃表现是判别器损失一路降到接近 0生成器损失却在原地踏步这时候你去看生成器输出大概率全是一个常数向量。原因通常是判别器能力太强生成器没机会学到东西。干预手段有三个按优先级排序第一把 hint_rate 降一降削弱判别器的信息优势第二调小判别器的隐藏层宽度或者往判别器加 Dropout降低它拟合速度第三生成器和判别器的学习率不要等比例我一般把判别器学习率设为生成器的 0.5 倍让它们之间保持一个“追赶但追不上”的节奏。3.4 一个让训练稳定下来的小习惯先训练判别器在每一个训练 step 里我会先更新判别器再更新生成器顺序上保持“判别器永远比生成器快半步”。这不是我拍脑袋想的而是原来把生成器放在前面训练生成器早期输出质量太差递进给判别器的全是垃圾样本判别器很容易学出一个“全盘否定”的决策边界后面怎么拉都拉不回来。先让判别器在某个 batch 上认清现状再让生成器针对性迷惑它这个对抗压力是持续有效的。4. 完整PyTorch代码落地从掩码构造到填补效果评估4.1 环境与依赖建议直接用 conda 建一个干净环境Python 3.8 以上即可PyTorch 1.10 以上都能跑我测试用的版本是 PyTorch 2.0CUDA 版本没有特殊要求CPU 也能跑只是 MNIST 全量训练会慢一些。基础依赖就四个torch、numpy、pandas、scikit-learn可视化用 matplotlib。4.2 随机缺失掩码构造训练 GAIN 需要一个“带缺失的数据集”。真实场景里缺失模式是数据自带的但为了验证效果通常会在完整数据集上人工构造缺失掩码。下面这段代码生成随机缺失比例下的二值掩码import torch import numpy as np def generate_mask(data, miss_rate0.2, seed42): torch.manual_seed(seed) batch_size, dim data.shape mask torch.rand(batch_size, dim) miss_rate return mask.float()注意这里每个位置独立缺失是典型的缺失完全随机模式。如果你的业务场景是某些整列缺失或者结构性缺失掩码生成逻辑要相应调整但后续网络部分完全不用改。4.3 生成器和判别器的完整定义网络结构我采用三层全连接隐藏维度取特征维度的 4 倍激活函数用 ReLU。生成器输出层不加激活靠外部 Clamp 限制范围。判别器输出层也不加激活配合 BCEWithLogits 计算损失代码上更稳定。import torch.nn as nn class Generator(nn.Module): def __init__(self, dim): super(Generator, self).__init__() self.model nn.Sequential( nn.Linear(dim * 2, dim * 4), nn.ReLU(), nn.Linear(dim * 4, dim * 4), nn.ReLU(), nn.Linear(dim * 4, dim), ) self.dim dim def forward(self, x_tilde, mask): inp torch.cat([x_tilde, mask], dim1) out self.model(inp) return torch.clamp(out, -10.0, 10.0) class Discriminator(nn.Module): def __init__(self, dim): super(Discriminator, self).__init__() self.model nn.Sequential( nn.Linear(dim * 2, dim * 4), nn.ReLU(), nn.Linear(dim * 4, dim * 4), nn.ReLU(), nn.Linear(dim * 4, dim), ) def forward(self, x_complete, hint): inp torch.cat([x_complete, hint], dim1) return self.model(inp)如果你的数据特征维度特别高比如基因表达数据有几千维隐藏层宽度不要动不动乘 8否则显存会顶不住而且训练非常容易过拟合。我一般遵循一个原则隐藏层不超过输入维度的 4 倍超过 1024 就封顶。4.4 训练循环核心代码训练循环是整个实现最需要抠细节的地方。下面给出完整的单 step 训练逻辑def train_one_step(generator, discriminator, data, mask, g_optim, d_optim, alpha10.0, hint_rate0.9): device data.device d_optim.zero_grad() # 用噪声填充缺失位置构造生成器输入 noise torch.rand_like(data).to(device) x_tilde mask * data (1 - mask) * noise # 生成器前向得到完整填补结果 x_hat generator(x_tilde, mask) x_complete mask * data (1 - mask) * x_hat # 构造 Hint 矩阵 hint_mask torch.rand_like(mask) hint_rate hint mask * hint_mask.float() 0.5 * (~hint_mask).float() # 判别器训练 pred_mask discriminator(x_complete.detach(), hint) loss_d nn.functional.binary_cross_entropy_with_logits(pred_mask, mask) loss_d.backward() d_optim.step() # 生成器训练 g_optim.zero_grad() pred_mask discriminator(x_complete, hint) ones torch.ones_like(mask) loss_g_adv nn.functional.binary_cross_entropy_with_logits(pred_mask, ones) loss_g_rec nn.functional.mse_loss(x_hat, data, reductionnone) loss_g_rec (loss_g_rec * mask).sum() / mask.sum() loss_g loss_g_adv alpha * loss_g_rec loss_g.backward() g_optim.step() return loss_d.item(), loss_g.item(), loss_g_adv.item(), loss_g_rec.item()这里有个很多人容易写错的地方训练判别器时传给判别器的 x_complete 要 .detach()切断梯度回传到生成器的路径。虽然即使不切断优化器只更新判别器参数生成器参数不会被误更新但梯度依然会流过生成器白白占用显存而且梯度累积状态会被污染长期跑会出一些奇怪现象。训练生成器时则不能 detach因为生成器的梯度必须通过判别器回传。4.5 评估指标与可视化训练结束后用一个独立构造缺失的测试集来评估。最直接的指标是 RMSE只在真实缺失位置上计算def evaluate(generator, data, mask): generator.eval() device data.device with torch.no_grad(): noise torch.rand_like(data).to(device) x_tilde mask * data (1 - mask) * noise x_hat generator(x_tilde, mask) x_complete mask * data (1 - mask) * x_hat true_values data[1 - mask 1] pred_values x_complete[1 - mask 1] rmse torch.sqrt(((true_values - pred_values) ** 2).mean()).item() return rmse如果数据是图像这类可可视化样本直接把 X_complete 的填补位置画成图像对比原图比任何数值指标都直观。我的建议是每个 epoch 保存一组可视化结果肉眼观察填补形状是否合理光看输出数字不行很多隐性错误要靠看才能发现。5. 实测对比与调参避坑数据分布决定填补上限5.1 MNIST上的填补效果从模糊到清晰我在 MNIST 上用 20% 随机缺失做了一轮完整测试训练 50 个 epochalpha 取 10hint_rate 取 0.9。前 10 个 epoch 生成的数字边缘模糊很多填补区域是均匀灰色到 20 个 epoch 左右轮廓出现笔画开始连贯30 个 epoch 之后填补位置已经能比较自然地和观察位置衔接起来。对比均值填补的输出均值填补在小面积缺失时看着还行一旦笔画的中间段缺失它填出来的完全是一坨灰雾因为它在所有缺失位置填的是同一个像素均值。5.2 和均值填补、MICE的量化对比为了不让结论停留在感觉层面我在同一个测试集上算了一组 RMSE 对比方法RMSE越低越好均值填补0.0183软插补/EM类方法0.0152MICEIterativeImputer0.0161GAINepoch200.0145GAINepoch500.0117数据是 0-1 标准化之后的 MNIST 特征所以 RMSE 数值看起来都不大。GAIN 在 50 个 epoch 时的优势已经很明确比均值填补低了接近 40%。更重要的是GAIN 填补后训练的线性分类器准确率比均值填补高出约 2 个百分点说明它保留的分布信息对下游任务有实际帮助。5.3 实际操作中踩过的几个深坑第一个坑是数据没有标准化就送进网络。GAIN 生成器的输出层默认不加激活如果原始数据量纲差异巨大生成器早期输出会直接爆到几百梯度瞬间消失模型彻底学不进去。后来我统一先做 Z-score 标准化训练结束后再把填补结果反标准化回原尺度。第二个坑是缺失比例和 hint_rate 的联动。当缺失率超过 40% 时hint_rate 保持 0.9 会导致判别器太强生成器崩溃。我摸索出来的规律是缺失率每升高 10 个百分点hint_rate 降 0.0560% 缺失率下用 0.75 左右比较合适。第三个坑是 batch size 太小导致判别器过拟合到局部模式。用 UCI 的小型表格数据时整个数据集才几百条我一开始用 16 的 batch size训练损失乱跳。后来改成一次全批量训练模型反而很快收敛。对小数据集建议全批量训练对大数据集再考虑 mini-batch。5.4 这套代码还能往哪个方向扩展GAIN 的适用范围远不止普通表格数据。我把同样的网络结构搬到单变量时序数据缺失填补上只需把输入改成滑动窗口切片效果比线性插值强很多。在真实业务数据上如果你明确知道缺失是“非随机缺失”也就是缺失本身带有信息可以考虑把缺失指示特征作为一个额外输入列喂给生成器让模型自动学习缺失模式和被填值之间的关系。另一个实用方向是用 GAIN 做数据增强里的“条件生成”比如只有部分特征被观测到时生成出完整的样本供下游模型训练。这套 PyTorch 代码的结构足够干净改改数据加载部分就能复用。最后再分享一个经验别把对抗训练的 epoch 数拍脑袋定死。我的做法是每隔 5 个 epoch 在验证集上评一次填补 RMSE连续 3 次不下降就提前停止。实战下来这个早停策略能帮我省掉大量训练时间也避免了后期过拟合导致填补质量变差。GAIN 的问题是调参但它调明白之后在缺失数据填补这个赛道上确实是能打的方案。本文还有配套的精品资源点击获取