使用256x256自定义数据集训练StyleGAN2出现TensorFlow维度不匹配错误
错误根源
你遇到的维度不匹配报错核心原因是训练配置和数据集分辨率不匹配:
你指定的--config=config-f是StyleGAN2官方预设的适配512x512分辨率训练的配置,内置的判别器、生成器结构默认按照512x512的输入输出尺寸设计,但你的自定义数据集分辨率为256x256,二者维度冲突触发报错。
另外你看到的tcmalloc: large alloc警告是config-f配置本身参数量大,单GPU训练时内存/显存占用过高的前兆,解决维度问题后也需要注意显存占用情况。
解决方案
可任选以下一种方法解决问题:
- 更换适配256x256分辨率的预设配置,将训练命令中的
--config=config-f修改为--config=config-e即可,config-e是官方预设的256x256分辨率训练专属配置,和你的数据集完全匹配。 - 如果你需要保留config-f的网络结构,在训练命令中添加
--resolution=256参数,显式指定训练分辨率,强制修改模型的输入输出维度适配你的256x256数据集。
针对Colab单GPU显存不足的问题,建议同时在训练命令中添加--batch-size=16(可根据实际运行情况调小数值),降低批量大小减少显存占用,避免后续训练出现OOM崩溃。
内容的提问来源于stack exchange,提问作者SEAN
相关产品推荐
相关产品推荐

