拓冰网站定制 
extreme_memorization 复现指南初始化尺度与极端记忆化实验详解ICLR 2021【免费下载链接】google-researchGoogle Research项目地址: https://gitcode.com/gh_mirrors/go/google-research本文是 Google Research 仓库中extreme_memorization目录的配套技术指南。该目录是 ICLR 2021 论文《Extreme Memorization via Scale of Initialization》的官方补充代码通过控制第一层权重 $w_1$ 的初始化标准差stddev、激活函数与损失函数系统性地验证过大的初始化尺度会导致网络对训练样本的极端记忆化这一现象并同时测量表征与梯度的类内对齐alignment指标。读完本文你将掌握三个数据集的 TFRecord 准备流程、全部命令行参数的含义与取值、论文第 3/4/5 节实验的逐条复现命令以及从 alignment.py、train.py 源码层面理解对齐度量与训练循环的底层实现。项目背景与代码结构该目录是论文《Extreme Memorization via Scale of Initialization》的补充代码核心研究问题是仅改变神经网络第一层权重的初始化尺度scale是否就能诱发对训练数据的极端记忆化。代码提供了一个最小可复现环境以 CIFAR-10 / CIFAR-100 / SVHN 三个图像分类数据集为基准在 2 层 MLP 与小型 ConvNet 上完成训练并把论文中关键的对齐alignment度量以 TensorBoard 标量的形式实时记录。仓库目录均位于extreme_memorization/下包含以下文件train.py训练/评估主入口使用 TensorFlow 2 eager 模式定义了全部命令行 Flag 与训练循环mlp.py2 层 MLP 模型支持自定义 $w_1$ 的随机正态初始化convnet.py论文第 5 节使用的卷积网络alignment.py类内对齐度量alignment的 O(n) 高效实现cifar10_dataset.py、cifar100_dataset.py、svhn_dataset.py三个数据集的tf.data.Dataset加载器requirements.txt依赖清单absl-py0.7.1、tensorflow2.1.0run.sh一键创建虚拟环境并安装依赖的脚本。环境准备数据集下载与依赖安装1. 下载并解压数据集论文实验使用三个 32×32 的经典图像数据集。官方提供的是已转换为 TFRecord 格式文件名含fingerprint标识的压缩包CIFAR-10https://www.dropbox.com/s/sstz260o6ad0ryn/cifar10.zip?dl0CIFAR-100https://www.dropbox.com/s/czun89ff9dhlc80/cifar100.zip?dl0SVHNhttps://www.dropbox.com/s/kf0p35wmp10j55x/svhn.zip?dl0解压后拷贝到/tmp/目录训练命令的默认输入路径即指向/tmpunzip cifar10.zip unzip cifar100.zip unzip svhn.zip cp -r cifar10 /tmp/ cp -r cifar100 /tmp/ cp -r svhn /tmp/2. 安装依赖pip install -r requirements.txt从 requirements.txt 可见依赖仅两个absl-py0.7.1提供 Flag 解析与日志与tensorflow2.1.0。仓库同时提供了 run.sh可一键完成虚拟环境创建与安装virtualenv -p python3 env source env/bin/activate pip install -r extreme_memorization/requirements.txt注意代码基于 TensorFlow 2.1.0 编写使用tensorflow.compat.v2导入与 eager 执行模式在更新版本环境中运行可能需要相应适配。3. TFRecord 数据格式说明从 cifar10_dataset.py 可以看出数据集中每条tf.Example包含三个字段特征名类型用途image/class/labelint64真实标签image/class/shuffled_labelint64打乱后的标签第 5 节 shuffled-labels 实验使用image/encodedstringPNG 编码的图像解码流程为decode_png→ 转 float32 →per_image_standardization逐图像标准化→ reshape 为长度为32*32*3 3072的一维向量cifar10_dataset.py。SVHN 数据集的原始字段名不同label与image由 svhn_dataset.py 在解析时统一映射为上述三个标准字段名。主要输入参数训练入口为python -m extreme_memorization.train通过 absl Flag 控制全部实验变量。README 列出的核心参数如下参数可选值 / 默认值说明--train_input_files无默认实际默认/tmp/cifar10/image_cifar10_fingerprint-train*训练 TFRecord 的输入通配符模式--test_input_files无默认实际默认/tmp/cifar10/image_cifar10_fingerprint-dev*测试 TFRecord 的输入通配符模式--datasetcifar10默认/cifar100/svhn数据集选择--model_typemlp默认/convnet模型架构2 层 MLP 或 ConvNet--activationrelu默认/sin/sigmoid激活函数仅对 MLP 生效--loss_functioncross_entropy默认/hinge/l2损失函数交叉熵、多类合页损失、平方损失--shuffled_labelsFalse默认是否使用打乱后的标签训练--custom_initFalse默认是否用自定义随机正态初始化器初始化 $w_1$仅对 MLP 生效--stddev0.001默认随机正态初始化的标准差仅在 MLP custom_inittrue时生效--num_units1024默认隐藏层单元数除上述参数外train.py 还定义了若干可调参数供精细化实验使用--train_epochs默认 2000训练轮数注释标明约 1000 轮起步--learning_rate默认 0.01优化器学习率优化器为tf.keras.optimizers.SGD--batch_size默认 256训练与评估的批大小多 GPU 下为全局批大小--log_interval默认 10每 N 个 batch 打印一次训练状态--output_dir默认/tmp/tensorflow/generalization/TensorBoard 摘要写入目录--model_dir默认/tmp/tensorflow/generalization/checkpoints/checkpoint 保存目录--no_gpu默认 False设为 True 时禁用 GPU--data_formatchannels_first/channels_last不指定时依据设备自动选择GPU 用channels_firstCPU 用channels_last见 train.py。可视化TensorBoard 中的对齐指标每次训练运行都会把表征与梯度的对齐度量写入 TF summary 标量可在 TensorBoard 中查看tensorboard --logdir/tmp/tensorflow/generalization结合源码训练过程中会记录以下几类标量train.pyloss/accuracy训练与测试eval子目录的损失与准确率representation_alignment/label每个类别在隐藏层表征上的类内对齐值由 mlp.py 与 convnet.py 中的plot_class_alignment写入grad_alignment/label/variable每个类别、每个变量含 logits上逐样本梯度的对齐值train.pyNorm/Grad/var与Norm/Var/var各变量的梯度范数与参数范数。其中对齐度量由 alignment.py 的compute_alignment计算先将向量按均值范数归一化再用公式(‖Σvᵢ‖² − Σ‖vᵢ‖²) / (n·(n−1))在 O(n) 时间内求得类内向量两两内积的平均值。该值越高表示同类别样本的表征或梯度方向越一致——这正是论文用于刻画记忆化程度的关键信号。复现第 3 节极端记忆化sin 激活 变化初始化尺度第 3 节实验使用sin 激活函数并将 $w_1$ 以随机正态分布初始化标准差在[0.001, 0.01, 0.1, 1.0, 10.0]范围内变化。以 CIFAR-10、stddev0.001为例python -m extreme_memorization.train \ --train_input_files/tmp/cifar10/image_cifar10_fingerprint-train* \ --test_input_files/tmp/cifar10/image_cifar10_fingerprint-dev* \ --activationsin \ --custom_inittrue \ --stddev0.001如何复现完整的 stddev 扫描将--stddev依次替换为0.01、0.1、1.0、10.0重复运行上述命令每次运行可配合--output_dir与--model_dir指向不同目录以便在 TensorBoard 中对比不同初始化尺度下的对齐曲线与 loss 收敛行为。从源码看custom_inittrue时 mlp.py 会以RandomNormal(mean0.0, stddevstddev)初始化隐藏层 Dense 的 kernel且use_biasFalse而 sin 激活在 train.py 中映射为tf.math.sin。注意激活函数仅在model_typemlp时生效。复现第 4 节为什么尺度会影响齐次激活ReLU 变化初始化尺度第 4 节将激活函数切换为ReLU齐次激活函数sin 的非齐次性质被排除后研究初始化尺度对 ReLU 网络记忆化的影响$w_1$ 仍以随机正态初始化stddev 在[0.001, 0.01, 0.1, 1.0, 10.0]范围内变化。ReLU 交叉熵stddev0.001CIFAR-10python -m extreme_memorization.train \ --train_input_files/tmp/cifar10/image_cifar10_fingerprint-train* \ --test_input_files/tmp/cifar10/image_cifar10_fingerprint-dev* \ --activationrelu \ --custom_inittrue \ --stddev0.001ReLU 合页损失hinge lossstddev0.001CIFAR-10python -m extreme_memorization.train \ --train_input_files/tmp/cifar10/image_cifar10_fingerprint-train* \ --test_input_files/tmp/cifar10/image_cifar10_fingerprint-dev* \ --activationrelu \ --custom_inittrue \ --loss_functionhinge \ --stddev0.001ReLU 平方损失squared lossstddev0.001CIFAR-10python -m extreme_memorization.train \ --train_input_files/tmp/cifar10/image_cifar10_fingerprint-train* \ --test_input_files/tmp/cifar10/image_cifar10_fingerprint-dev* \ --activationrelu \ --custom_inittrue \ --loss_functionl2 \ --stddev0.001三种损失函数的实现位于 train.pycross_entropysparse_softmax_cross_entropy_with_logitshinge多类合页损失实现为Σ_j max(0, logits_j − logits_{label} 1) − 1见hinge_loss与辅助的gather_2d其中gather_2d用于按批索引取标签对应的 logitsl2对 one-hot 标签向量的均方误差。通过对比三种损失函数在同一初始化尺度下的对齐与过拟合曲线可以检验记忆化由梯度方向对齐驱动这一结论是否对损失函数选择稳健。复现第 5 节对齐现象是否具有更广的普适性第 5 节将激活函数保持为 ReLU但所有变量改用 glorot uniform 初始化器即不设custom_initmlp.py 的默认分支使用 Keras Dense 默认的 glorot uniform 初始化用于验证对齐现象不依赖于特殊初始化。在 CIFAR-10 上训练 2 层 MLPpython -m extreme_memorization.train \ --train_input_files/tmp/cifar10/image_cifar10_fingerprint-train* \ --test_input_files/tmp/cifar10/image_cifar10_fingerprint-dev* \ --activationrelu在 CIFAR-10 上使用打乱标签shuffled labels训练python -m extreme_memorization.train \ --train_input_files/tmp/cifar10/image_cifar10_fingerprint-train* \ --test_input_files/tmp/cifar10/image_cifar10_fingerprint-dev* \ --activationrelu \ --shuffled_labelstrue当--shuffled_labelstrue时train.py 的get_image_labels会读取 TFRecord 中预先生成的image/class/shuffled_label字段代替真实标签。这是经典的标签打乱对照实验模型只能依靠纯记忆拟合随机标签用于观察记忆化达到极端程度时对齐指标的行为。在 CIFAR-10 上训练 ConvNetpython -m extreme_memorization.train \ --train_input_files/tmp/cifar10/image_cifar10_fingerprint-train* \ --test_input_files/tmp/cifar10/image_cifar10_fingerprint-dev* \ --model_typeconvnetConvNet 架构定义在 convnet.pyConv2D(32, 5, same) → MaxPool(3,3,stride 2) → Conv2D(64, 5, valid) → MaxPool(3,3,stride 2) → Flatten → Dense(1024, relu) → Dense(num_labels)在隐藏层之后同样会记录representation_alignment。该实验用于说明对齐驱动的记忆化不仅存在于 MLP在卷积网络上同样可以观察。训练循环内部机制源码视角若要深入理解每次运行到底发生了什么可以关注 train.py 的训练函数持久化梯度带使用tf.GradientTape(persistentTrue)同时监视输入图像与 logits逐样本梯度收集对 batch 内每个样本单独计算 loss 对 logits 与各变量的梯度并按标签分组per_label_grads逐类梯度对齐对每个标签最多 10 类下各变量的逐样本梯度计算对齐值并写入grad_alignment/label/var标量梯度更新对整体 loss 求梯度后用 SGD 更新变量同时记录各变量梯度范数与参数范数checkpoint每个 epoch 结束后在--model_dir保存 checkpoint并支持从已有 checkpoint 恢复train.py便于中断后续跑。设备与数据格式方面run_eager会优先使用 GPUchannels_first数据格式无 GPU 或指定--no_gpu时回退 CPUchannels_last该选择也影响 convnet.py 中对输入的转置处理。常见问题与注意事项数据集路径README 中的示例命令默认数据位于/tmp/cifar10/切换数据集时将--dataset改为cifar100或svhn并将--train_input_files/--test_input_files指向对应目录下的fingerprint-train*与fingerprint-dev*通配符。custom_init仅对 MLP 生效ConvNet 不支持自定义初始化stddev参数也仅在 MLP custom_inittrue时才有意义。激活函数仅对 MLP 生效--activationsin/sigmoid/relu不会影响 ConvNet其内部固定使用 ReLU。CIFAR-100 的对齐记录开销代码注释指出逐类梯度对齐默认只统计前 10 个标签train.py因为对 100 类全部统计会带来显著计算开销。长训练时长默认--train_epochs2000单次实验耗时较长可先用--train_epochs调小做快速冒烟测试确认数据路径与参数无误后再跑完整实验。Python 模块运行方式命令通过python -m extreme_memorization.train执行需在仓库根目录下运行确保extreme_memorization可作为顶层包被导入。【免费下载链接】google-researchGoogle Research项目地址: https://gitcode.com/gh_mirrors/go/google-research创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 extreme_memorization 复现指南：初始化尺度与“极端记忆化“实验详解（ICLR 2021）
extreme_memorization 复现指南初始化尺度与极端记忆化实验详解ICLR 2021【免费下载链接】google-researchGoogle Research项目地址: https://gitcode.com/gh_mirrors/go/google-research本文是 Google Research 仓库中extreme_memorization目录的配套技术指南。该目录是 ICLR 2021 论文《Extreme Memorization via Scale of Initialization》的官方补充代码通过控制第一层权重 $w_1$ 的初始化标准差stddev、激活函数与损失函数系统性地验证过大的初始化尺度会导致网络对训练样本的极端记忆化这一现象并同时测量表征与梯度的类内对齐alignment指标。读完本文你将掌握三个数据集的 TFRecord 准备流程、全部命令行参数的含义与取值、论文第 3/4/5 节实验的逐条复现命令以及从 alignment.py、train.py 源码层面理解对齐度量与训练循环的底层实现。项目背景与代码结构该目录是论文《Extreme Memorization via Scale of Initialization》的补充代码核心研究问题是仅改变神经网络第一层权重的初始化尺度scale是否就能诱发对训练数据的极端记忆化。代码提供了一个最小可复现环境以 CIFAR-10 / CIFAR-100 / SVHN 三个图像分类数据集为基准在 2 层 MLP 与小型 ConvNet 上完成训练并把论文中关键的对齐alignment度量以 TensorBoard 标量的形式实时记录。仓库目录均位于extreme_memorization/下包含以下文件train.py训练/评估主入口使用 TensorFlow 2 eager 模式定义了全部命令行 Flag 与训练循环mlp.py2 层 MLP 模型支持自定义 $w_1$ 的随机正态初始化convnet.py论文第 5 节使用的卷积网络alignment.py类内对齐度量alignment的 O(n) 高效实现cifar10_dataset.py、cifar100_dataset.py、svhn_dataset.py三个数据集的tf.data.Dataset加载器requirements.txt依赖清单absl-py0.7.1、tensorflow2.1.0run.sh一键创建虚拟环境并安装依赖的脚本。环境准备数据集下载与依赖安装1. 下载并解压数据集论文实验使用三个 32×32 的经典图像数据集。官方提供的是已转换为 TFRecord 格式文件名含fingerprint标识的压缩包CIFAR-10https://www.dropbox.com/s/sstz260o6ad0ryn/cifar10.zip?dl0CIFAR-100https://www.dropbox.com/s/czun89ff9dhlc80/cifar100.zip?dl0SVHNhttps://www.dropbox.com/s/kf0p35wmp10j55x/svhn.zip?dl0解压后拷贝到/tmp/目录训练命令的默认输入路径即指向/tmpunzip cifar10.zip unzip cifar100.zip unzip svhn.zip cp -r cifar10 /tmp/ cp -r cifar100 /tmp/ cp -r svhn /tmp/2. 安装依赖pip install -r requirements.txt从 requirements.txt 可见依赖仅两个absl-py0.7.1提供 Flag 解析与日志与tensorflow2.1.0。仓库同时提供了 run.sh可一键完成虚拟环境创建与安装virtualenv -p python3 env source env/bin/activate pip install -r extreme_memorization/requirements.txt注意代码基于 TensorFlow 2.1.0 编写使用tensorflow.compat.v2导入与 eager 执行模式在更新版本环境中运行可能需要相应适配。3. TFRecord 数据格式说明从 cifar10_dataset.py 可以看出数据集中每条tf.Example包含三个字段特征名类型用途image/class/labelint64真实标签image/class/shuffled_labelint64打乱后的标签第 5 节 shuffled-labels 实验使用image/encodedstringPNG 编码的图像解码流程为decode_png→ 转 float32 →per_image_standardization逐图像标准化→ reshape 为长度为32*32*3 3072的一维向量cifar10_dataset.py。SVHN 数据集的原始字段名不同label与image由 svhn_dataset.py 在解析时统一映射为上述三个标准字段名。主要输入参数训练入口为python -m extreme_memorization.train通过 absl Flag 控制全部实验变量。README 列出的核心参数如下参数可选值 / 默认值说明--train_input_files无默认实际默认/tmp/cifar10/image_cifar10_fingerprint-train*训练 TFRecord 的输入通配符模式--test_input_files无默认实际默认/tmp/cifar10/image_cifar10_fingerprint-dev*测试 TFRecord 的输入通配符模式--datasetcifar10默认/cifar100/svhn数据集选择--model_typemlp默认/convnet模型架构2 层 MLP 或 ConvNet--activationrelu默认/sin/sigmoid激活函数仅对 MLP 生效--loss_functioncross_entropy默认/hinge/l2损失函数交叉熵、多类合页损失、平方损失--shuffled_labelsFalse默认是否使用打乱后的标签训练--custom_initFalse默认是否用自定义随机正态初始化器初始化 $w_1$仅对 MLP 生效--stddev0.001默认随机正态初始化的标准差仅在 MLP custom_inittrue时生效--num_units1024默认隐藏层单元数除上述参数外train.py 还定义了若干可调参数供精细化实验使用--train_epochs默认 2000训练轮数注释标明约 1000 轮起步--learning_rate默认 0.01优化器学习率优化器为tf.keras.optimizers.SGD--batch_size默认 256训练与评估的批大小多 GPU 下为全局批大小--log_interval默认 10每 N 个 batch 打印一次训练状态--output_dir默认/tmp/tensorflow/generalization/TensorBoard 摘要写入目录--model_dir默认/tmp/tensorflow/generalization/checkpoints/checkpoint 保存目录--no_gpu默认 False设为 True 时禁用 GPU--data_formatchannels_first/channels_last不指定时依据设备自动选择GPU 用channels_firstCPU 用channels_last见 train.py。可视化TensorBoard 中的对齐指标每次训练运行都会把表征与梯度的对齐度量写入 TF summary 标量可在 TensorBoard 中查看tensorboard --logdir/tmp/tensorflow/generalization结合源码训练过程中会记录以下几类标量train.pyloss/accuracy训练与测试eval子目录的损失与准确率representation_alignment/label每个类别在隐藏层表征上的类内对齐值由 mlp.py 与 convnet.py 中的plot_class_alignment写入grad_alignment/label/variable每个类别、每个变量含 logits上逐样本梯度的对齐值train.pyNorm/Grad/var与Norm/Var/var各变量的梯度范数与参数范数。其中对齐度量由 alignment.py 的compute_alignment计算先将向量按均值范数归一化再用公式(‖Σvᵢ‖² − Σ‖vᵢ‖²) / (n·(n−1))在 O(n) 时间内求得类内向量两两内积的平均值。该值越高表示同类别样本的表征或梯度方向越一致——这正是论文用于刻画记忆化程度的关键信号。复现第 3 节极端记忆化sin 激活 变化初始化尺度第 3 节实验使用sin 激活函数并将 $w_1$ 以随机正态分布初始化标准差在[0.001, 0.01, 0.1, 1.0, 10.0]范围内变化。以 CIFAR-10、stddev0.001为例python -m extreme_memorization.train \ --train_input_files/tmp/cifar10/image_cifar10_fingerprint-train* \ --test_input_files/tmp/cifar10/image_cifar10_fingerprint-dev* \ --activationsin \ --custom_inittrue \ --stddev0.001如何复现完整的 stddev 扫描将--stddev依次替换为0.01、0.1、1.0、10.0重复运行上述命令每次运行可配合--output_dir与--model_dir指向不同目录以便在 TensorBoard 中对比不同初始化尺度下的对齐曲线与 loss 收敛行为。从源码看custom_inittrue时 mlp.py 会以RandomNormal(mean0.0, stddevstddev)初始化隐藏层 Dense 的 kernel且use_biasFalse而 sin 激活在 train.py 中映射为tf.math.sin。注意激活函数仅在model_typemlp时生效。复现第 4 节为什么尺度会影响齐次激活ReLU 变化初始化尺度第 4 节将激活函数切换为ReLU齐次激活函数sin 的非齐次性质被排除后研究初始化尺度对 ReLU 网络记忆化的影响$w_1$ 仍以随机正态初始化stddev 在[0.001, 0.01, 0.1, 1.0, 10.0]范围内变化。ReLU 交叉熵stddev0.001CIFAR-10python -m extreme_memorization.train \ --train_input_files/tmp/cifar10/image_cifar10_fingerprint-train* \ --test_input_files/tmp/cifar10/image_cifar10_fingerprint-dev* \ --activationrelu \ --custom_inittrue \ --stddev0.001ReLU 合页损失hinge lossstddev0.001CIFAR-10python -m extreme_memorization.train \ --train_input_files/tmp/cifar10/image_cifar10_fingerprint-train* \ --test_input_files/tmp/cifar10/image_cifar10_fingerprint-dev* \ --activationrelu \ --custom_inittrue \ --loss_functionhinge \ --stddev0.001ReLU 平方损失squared lossstddev0.001CIFAR-10python -m extreme_memorization.train \ --train_input_files/tmp/cifar10/image_cifar10_fingerprint-train* \ --test_input_files/tmp/cifar10/image_cifar10_fingerprint-dev* \ --activationrelu \ --custom_inittrue \ --loss_functionl2 \ --stddev0.001三种损失函数的实现位于 train.pycross_entropysparse_softmax_cross_entropy_with_logitshinge多类合页损失实现为Σ_j max(0, logits_j − logits_{label} 1) − 1见hinge_loss与辅助的gather_2d其中gather_2d用于按批索引取标签对应的 logitsl2对 one-hot 标签向量的均方误差。通过对比三种损失函数在同一初始化尺度下的对齐与过拟合曲线可以检验记忆化由梯度方向对齐驱动这一结论是否对损失函数选择稳健。复现第 5 节对齐现象是否具有更广的普适性第 5 节将激活函数保持为 ReLU但所有变量改用 glorot uniform 初始化器即不设custom_initmlp.py 的默认分支使用 Keras Dense 默认的 glorot uniform 初始化用于验证对齐现象不依赖于特殊初始化。在 CIFAR-10 上训练 2 层 MLPpython -m extreme_memorization.train \ --train_input_files/tmp/cifar10/image_cifar10_fingerprint-train* \ --test_input_files/tmp/cifar10/image_cifar10_fingerprint-dev* \ --activationrelu在 CIFAR-10 上使用打乱标签shuffled labels训练python -m extreme_memorization.train \ --train_input_files/tmp/cifar10/image_cifar10_fingerprint-train* \ --test_input_files/tmp/cifar10/image_cifar10_fingerprint-dev* \ --activationrelu \ --shuffled_labelstrue当--shuffled_labelstrue时train.py 的get_image_labels会读取 TFRecord 中预先生成的image/class/shuffled_label字段代替真实标签。这是经典的标签打乱对照实验模型只能依靠纯记忆拟合随机标签用于观察记忆化达到极端程度时对齐指标的行为。在 CIFAR-10 上训练 ConvNetpython -m extreme_memorization.train \ --train_input_files/tmp/cifar10/image_cifar10_fingerprint-train* \ --test_input_files/tmp/cifar10/image_cifar10_fingerprint-dev* \ --model_typeconvnetConvNet 架构定义在 convnet.pyConv2D(32, 5, same) → MaxPool(3,3,stride 2) → Conv2D(64, 5, valid) → MaxPool(3,3,stride 2) → Flatten → Dense(1024, relu) → Dense(num_labels)在隐藏层之后同样会记录representation_alignment。该实验用于说明对齐驱动的记忆化不仅存在于 MLP在卷积网络上同样可以观察。训练循环内部机制源码视角若要深入理解每次运行到底发生了什么可以关注 train.py 的训练函数持久化梯度带使用tf.GradientTape(persistentTrue)同时监视输入图像与 logits逐样本梯度收集对 batch 内每个样本单独计算 loss 对 logits 与各变量的梯度并按标签分组per_label_grads逐类梯度对齐对每个标签最多 10 类下各变量的逐样本梯度计算对齐值并写入grad_alignment/label/var标量梯度更新对整体 loss 求梯度后用 SGD 更新变量同时记录各变量梯度范数与参数范数checkpoint每个 epoch 结束后在--model_dir保存 checkpoint并支持从已有 checkpoint 恢复train.py便于中断后续跑。设备与数据格式方面run_eager会优先使用 GPUchannels_first数据格式无 GPU 或指定--no_gpu时回退 CPUchannels_last该选择也影响 convnet.py 中对输入的转置处理。常见问题与注意事项数据集路径README 中的示例命令默认数据位于/tmp/cifar10/切换数据集时将--dataset改为cifar100或svhn并将--train_input_files/--test_input_files指向对应目录下的fingerprint-train*与fingerprint-dev*通配符。custom_init仅对 MLP 生效ConvNet 不支持自定义初始化stddev参数也仅在 MLP custom_inittrue时才有意义。激活函数仅对 MLP 生效--activationsin/sigmoid/relu不会影响 ConvNet其内部固定使用 ReLU。CIFAR-100 的对齐记录开销代码注释指出逐类梯度对齐默认只统计前 10 个标签train.py因为对 100 类全部统计会带来显著计算开销。长训练时长默认--train_epochs2000单次实验耗时较长可先用--train_epochs调小做快速冒烟测试确认数据路径与参数无误后再跑完整实验。Python 模块运行方式命令通过python -m extreme_memorization.train执行需在仓库根目录下运行确保extreme_memorization可作为顶层包被导入。【免费下载链接】google-researchGoogle Research项目地址: https://gitcode.com/gh_mirrors/go/google-research创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考