怎样降低GeneFace显存占用

GeneFace（以及 GeneFace++）在推理/训练时显存占用主要来自 面部渲染（一般基于 NeRF / 3DMM 渲染）、音频特征提取模型（如 HuBERT / wav2vec） 以及 中间张量缓存。下面按「实用程度」给你一套降显存方案。

一、推理阶段（最常用，优先做）

1. 降低渲染分辨率

这是最立竿见影的方法。

# 常见参数
img_size = 256  # 改成 128 甚至 96

• 256 → 128：显存通常能降 40%~60%
• 人脸视频对 128 分辨率通常仍可接受

2. 减少 batch size（推理也有效）

即使推理，很多实现内部仍有 batch：

batch_size = 1

如果是视频生成：

• 一次只处理 1 段短音频
• 或分片生成再拼接

3. 关掉不必要的模型

GeneFace 常见组合：

• 音频编码器（HuBERT）
• 运动模型
• 渲染模型

如果你只做推理：

use_hubert = False  # 用预提取特征

提前提取音频特征并保存，避免常驻显存。

4. 使用 FP16 / BF16

在支持的设备上：

model.half()
audio_feature = audio_feature.half()

注意：

• 渲染部分对精度较敏感
• 若出现异常，可只给音频模型用 FP16

5. 启用 CPU offload（低显存神器）

PyTorch 示例：

model.to("cpu")  # 不用时
model.to("cuda") # 用时

或使用：

• accelerate
• torch.cuda.empty_cache() 定期清理

二、训练阶段（更吃显存）

6. 减小 NeRF / 渲染网络规模

GeneFace 常基于：

• MLP 层数
• hidden dim

hidden_dim = 128  # 原 256
n_layers = 4      # 原 8

会牺牲一点质量，但显存下降明显。

7. 梯度检查点（Gradient Checkpointing）

model.gradient_checkpointing_enable()

用时间换显存，训练速度略慢。

8. 只训练一部分模块

例如：

• 冻结音频编码器
• 冻结 3DMM 参数
• 只训渲染网络

for p in audio_encoder.parameters():
p.requires_grad = False

三、系统与环境层面

9. 使用更小 CUDA / PyTorch 版本

• PyTorch 2.x + CUDA 11.8 通常比老版本更省显存
• 避免同时开多个 Python 进程

10. 显存清理习惯

import torch
torch.cuda.empty_cache()

在每一段生成后调用。

四、推荐组合（实用方案）

6G 显存（如 2060 / 3060）

• img_size=128
• batch=1
• FP16
• 预提取音频特征
• 分片推理

4G 显存

• img_size=96
• CPU 提取音频
• 渲染模型常驻，其余 CPU offload

如果你愿意，可以告诉我：

• 你是 推理还是训练
• 显存大小（如 6G / 8G / 24G）
• 用的是 GeneFace 还是 GeneFace++

我可以直接帮你改对应配置文件或代码。