从环境到出图:AMD显卡上用kohya_ss训练AI绘画模型的完整实操指南
【免费下载链接】kohya_ss项目地址: https://gitcode.com/GitHub_Trending/ko/kohya_ss
kohya_ss 是一套在 AMD显卡 上完成 AI绘画模型训练 的开源工具,基于 ROCm 技术栈支持 LoRA 微调、DreamBooth 与 SDXL 完整训练,并配有 Gradio 图形界面,覆盖数据准备到模型产出的全过程。本文按"装好环境 → 跑出第一个模型 → 处理数据 → 调优 → 排障"的顺序,带你把这条链路完整走一遍。
装好ROCm训练环境:4步让AMD显卡被PyTorch识别
完成这一步后,你的 Python 环境里会有一个能直接调用 AMD GPU 的 PyTorch,后续所有训练命令都依赖它。
先确认系统满足三个前提:
- Ubuntu 20.04 / 22.04 LTS
- ROCm 6.3 或更高版本
- Python 3.10 或 3.11
然后按顺序执行:
# 1. 安装 ROCm 驱动组件 sudo apt update sudo apt install rocm-hip-sdk # 2. 获取 kohya_ss git clone https://gitcode.com/GitHub_Trending/ko/kohya_ss cd kohya_ss # 3. 安装 AMD 专用依赖(torch 2.7.1+rocm6.3、tensorflow-rocm、onnxruntime-rocm) pip install -r requirements_linux_rocm.txt # 4. 验证:版本正常输出且第二行为 True 即成功 python -c "import torch; print(torch.__version__); print(torch.cuda.is_available())"其中 requirements_linux_rocm.txt 指向 ROCm 6.3 的 PyTorch 官方索引,装错普通 CUDA 版 torch 是后续大部分问题的根源,验证命令的True就是最直接的确认信号。
第一次跑通训练:10张图片起步的最小配置
做完这一步,你将拥有第一个训练完成的模型文件和一批样本图,并建立起对三种训练模式耗时、体积的预期。
先用 test/img/ 里自带的小数据集试跑:约 10-20 张 512x512 的图片即可,不必等数据量充足才开始。
三种训练方式的量级对比:
| 训练模式 | 典型耗时 | 模型体积 | 适用情况 |
|---|---|---|---|
| 完整模型训练 | 8-12 小时 | 2-7GB | 数据量大、全新风格 |
| DreamBooth | 3-5 小时 | 2-4GB | 特定对象或人物 |
| LoRA 微调 | 1-2 小时 | 10-100MB | 少量数据、风格迁移 |
启动界面:
python kohya_gui.py在 Gradio 界面中选择训练模式后,新手建议这样起步(以 RX 7900 XTX 为例,其他显卡见下文):
| 参数 | 建议值 | 说明 |
|---|---|---|
| batch_size | 4 | 按显卡型号调整 |
| fp16 | 开启 | 降低显存占用 |
| learning_rate | 0.0003 | 可用区间 0.0001~0.0005 |
| rank | 4~16 | 越大模型容量越大 |
| 训练步数 | 500~2000 | 随数据量增减 |
训练数据工具:给图片补描述、分组和压缩
这一步的产出是一个"每张图有 caption 文件、按尺寸分好组、格式统一"的干净数据集,训练时不再因为缺描述或 IO 慢而卡住。
仓库 tools/ 目录里可以直接用:
# 为目录内图片自动生成文本描述 python tools/caption.py --input_dir=你的图片目录 # 按尺寸相似性把图片分组到不同文件夹 python tools/group_images.py --input_dir=你的图片目录此外convert_images_to_webp.py、convert_images_to_hq_jpg.py负责批量格式转换,crop_images_to_n_buckets.py负责裁切到分桶尺寸。kohya_gui/ 下还封装了图形化的分组、手动标注等入口,不想敲命令可以从界面走。
显存告急先调这3个参数:batch_size、检查点、FP16
做完这一步,你会按自己的显卡型号得到一组"能跑稳"的参数组合,而不是盲目试错。
三个旋钮的作用:
- batch_size:最直接的控制项,显存不够时优先降到 1 或 2。
- gradient_checkpointing:牺牲约 20% 的训练速度,换 30-50% 的显存节省,适合高分辨率图或大模型。
- fp16 混合精度:显存占用可接近减半,对出图质量影响很小。
按型号给的建议值:
| 显卡 | batch_size | 配套动作 |
|---|---|---|
| RX 7900 XTX | 4-8 | 配合梯度检查点 |
| RX 7800 XT | 2-4 | 开启 FP16 |
| RX 7700 XT | 1-2 | 同时降低图片分辨率 |
速度方面还有三个低成本动作:训练前把图片全部预处理完;数据集放在 SSD 上;按 CPU 核心数设置数据加载 workers。
四类高频故障:症状对照排查
每个问题都给出"看到什么 → 做什么"的对应关系。
症状 1:训练中途中断,提示显存不足
- batch_size 降到 1 或 2
- 开启 gradient_checkpointing
- 分辨率从 512x512 降到 384x384
- 拿 test/config/ 下的 dataset.toml 系列配置做参照修改
症状 2:启动即报 hipErrorNoBinaryForGpu
- 升级 ROCm 到 6.3+:
sudo apt upgrade rocm-hip-sdk - 运行
rocminfo确认驱动与设备状态 - 重启系统,确保新驱动完全加载
症状 3:导入 TensorFlow 报版本兼容错误
- 先确认 Python 版本与 requirements_linux_rocm.txt 的适配说明一致
- 清理旧包:
pip uninstall tensorflow tensorflow-rocm - 重新执行
pip install -r requirements_linux_rocm.txt
症状 4:GPU 利用率低、训练速度异常慢
rocminfo | grep "GPU"确认 ROCm 是否识别到卡- 确认装的是
+rocm6.3后缀的 PyTorch,而不是 CUDA 版 - 训练时用
rocm-smi实时观察显存与利用率
进阶与扩展:多卡、预设和内置监控
多 GPU 训练:在 GUI 中启用多卡选项,把总 batch_size 分摊到每张卡;显存吃紧时改用梯度累积,用更小的单卡 batch 模拟更大的等效 batch。
预设定制:presets/ 里全是可直接启用的 JSON 参数集。新手从SDXL - LoRA AI_characters standard v1.0.json入手看默认项,进阶后对照SDXL - LoRA kudou-reira prodigy v4.0.json理解优化器与学习率的配合,再把自己的成功参数存成模板。
训练监控:kohya_ss 内置实时 Loss 曲线、每 N 步自动生成的样本预览,以及显存使用情况,训练中途即可判断参数是否合理。
更多资料入口:docs/ 收录分平台的安装与训练文档(含 Linux/Windows 的 pip、uv 方案),test/ 提供完整的示例图片与配置模板,适合改完配置后快速回归验证。
【免费下载链接】kohya_ss项目地址: https://gitcode.com/GitHub_Trending/ko/kohya_ss
创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考