【Bug已解决】Fix incorrect batch handling in _prepare_image_ids usage in train_dreambooth_lora_flux2_img2

【Bug已解决】Fix incorrect batch handling in _prepare_image_ids usage in train_dreambooth_lora_flux2_img2img.py 解决方案一、现象长什么样train_dreambooth_lora_flux2_img2img.py是给 FLUX.2 做 dreambooth LoRA img2img 的训练脚本。当用户用batch size 1训练时结果不对劲loss 震荡、生成的 LoRA 在不同样本间串味或训练中途shape mismatchaccelerate launch train_dreambooth_lora_flux2_img2img.py \ --train_batch_size 4 --instance_data_dir ./imgs ...报错RuntimeError The expanded size of the tensor (4096) must match the existing size (1024) at non-singleton dimension 1或者训练能跑但质量差# batch1 没问题batch4 时 image ids 明显错位位置编码错定位到问题在_prepare_image_ids这个函数——它负责给每个训练样本生成「图像位置 id」用于 FLUX.2 的二维/三维旋转位置编码但在 batch 维度上处理错了要么把整个 batch 当成一张大图去算 id要么把单样本的 id 错误地 broadcast 到 batch 上。现象总结_prepare_image_ids在计算图像位置 id 时没正确处理 batch 维度——batch 内每个样本应有自己独立的图像 id基于各自的高宽但脚本把 batch 当单样本处理或错误共享 id导致位置编码在 batch1 时错位loss 异常或 shape 报错。二、背景FLUX.2和 FLUX 类似的图像位置 id 是根据图像的高、宽、以及 patch 切分算出来的每个 patch 有(h_id, w_id)二维坐标展平成一维 id 序列。这个 id 序列完全取决于单张图的高宽。训练脚本里_prepare_image_ids的职责是对任意一张图给定height,width,patch_size,axes_dim算出它的 id 张量。问题在于 batch 处理正确对 batch 中每个样本独立调用_prepare_image_ids得到[B, N_ids]每个样本自己的 id错误脚本现状_prepare_image_ids被喂入了整个 batch 的高宽或只用了 batch 里第 0 个样本的高宽于是算出的 id 只对应一张图的大小再 broadcast 到 batch 上时非第 0 个样本的位置 id 全错。当 batch 内各样本高宽一致时错误可能「悄悄」存在id 形状对但语义重复当高宽不一致时直接 shape mismatch。三、根因根因两点_prepare_image_ids把 batch 当单样本函数内部用height/width算 id但调用处传入的是 batch 级张量或只用了第 0 样本的高宽导致 batch 内样本共享或错用同一组 id。缺少「按样本独立计算 正确堆叠」的逻辑正确做法是对for i in range(B): ids[i] _prepare_image_ids(h[i], w[i])脚本却做了ids _prepare_image_ids(batch_h, batch_w)这种批量误用。本质图像位置 id 是「每样本独立」的量但训练脚本在 batch 维度上错误地把它当「全局共享/批量广播」破坏了batch1时的位置编码正确性。四、最小可运行复现用标准库复现「batch 内共享单样本 id 导致错位」import torch def prepare_image_ids(height, width, patch2): # 按单张图高宽算二维 id h_ids torch.arange(height // patch) w_ids torch.arange(width // patch) grid torch.stack(torch.meshgrid(h_ids, w_ids, indexingij), -1) return grid.reshape(-1, 2) # [N_ids, 2] def buggy_batch_call(batch_h, batch_w): # 错误只用了 batch 第 0 个样本的高宽broadcast 到整个 batch h0, w0 batch_h[0].item(), batch_w[0].item() ids prepare_image_ids(h0, w0) return ids.unsqueeze(0).expand(len(batch_h), -1, -1) # 所有样本共享第 0 个的 id def correct_batch_call(batch_h, batch_w): return torch.stack([prepare_image_ids(h.item(), w.item()) for h, w in zip(batch_h, batch_w)]) bh torch.tensor([64, 64, 96, 64]) # batch 内样本高不同 bw torch.tensor([64, 64, 64, 96]) buggy buggy_batch_call(bh, bw) correct correct_batch_call(bh, bw) print(buggy[2] shape:, buggy[2].shape, (基于 h64 算的错)) print(correct[2] shape:, correct[2].shape, (基于 h96 算的对)) assert buggy[2].shape ! correct[2].shape # 第 3 个样本 id 错位复现「shape mismatch」当bh里某个样本高宽和第 0 个差很多broadcast 出的 id 长度与 transformer 期望的 token 数不一致直接 RuntimeErr。五、解决方案第一层最小直接修复最小修复让_prepare_image_ids在 batch 维度按样本独立计算并正确堆叠调用处不再批量误用import torch def prepare_image_ids(height, width, patch_size2, axes_dim(16, 16)): 单样本返回 [N_patches, 2] 的 (h_id, w_id)。 h, w height // patch_size, width // patch_size h_ids torch.arange(h) w_ids torch.arange(w) grid torch.stack(torch.meshgrid(h_ids, w_ids, indexingij), -1) return grid.reshape(-1, 2) def prepare_image_ids_batch(batch_height, batch_width, patch_size2, axes_dim(16, 16)): 修复每个样本独立算 id再堆叠成 [B, N_i, 2]。 ids [ prepare_image_ids(h.item(), w.item(), patch_size, axes_dim) for h, w in zip(batch_height, batch_width) ] # 若 batch 内尺寸一致可 stack不一致需 pad这里假定一致训练通常固定尺寸 return torch.stack(ids, dim0)训练循环里改用prepare_image_ids_batch(batch_h, batch_w)每个样本拿到自己正确的位置 id。六、解决方案第二层结构性改进把「图像位置 id 的 batch 处理契约」收敛成一个 dataclass 单一真源防止再被批量误用from dataclasses import dataclass, field from typing import List, Tuple dataclass(frozenTrue) class ImageIdsBatchPolicy: _prepare_image_ids 的 batch 处理单一真源。 # 位置 id 是否每样本独立 per_sample_independent: bool True # id 维度h_id, w_id 或加 t id_dim: int 2 # patch 切分尺寸 patch_size: int 2 # 是否允许 batch 内尺寸不一致不一致需 pad allow_variable_size: bool False # 轴向维度FLUX 的 axes_dim axes_dim: Tuple[int, int] (16, 16) def prepare(self, batch_height, batch_width): if not self.per_sample_independent: raise ValueError(image ids 必须每样本独立计算) ids [self._single(h.item(), w.item()) for h, w in zip(batch_height, batch_width)] if self.allow_variable_size: ids self._pad_to_max(ids) return torch.stack(ids, dim0) def _single(self, h, w): hh, ww h // self.patch_size, w // self.patch_size grid torch.stack(torch.meshgrid(torch.arange(hh), torch.arange(ww), indexingij), -1) return grid.reshape(-1, self.id_dim) def _pad_to_max(self, ids): max_n max(i.shape[0] for i in ids) return [torch.cat([i, i.new_zeros(max_n - i.shape[0], self.id_dim)], 0) for i in ids] def validate_batch(self, batch_height, batch_width, out) - List[str]: problems [] if out.shape[0] ! len(batch_height): problems.append(输出 batch 维度与输入不一致) if not self.allow_variable_size: if any(h ! batch_height[0] or w ! batch_width[0] for h, w in zip(batch_height, batch_width)): problems.append(尺寸不一致但 allow_variable_sizeFalse) return problems训练脚本只用policy.prepare(batch_h, batch_w)validate_batch在训练前校验。七、解决方案第三层断言 / CI 守护用 pytest 把「batch 内每样本独立 id 形状一致 尺寸校验」固化成回归import torch import pytest from mylib.image_ids_batch import ImageIdsBatchPolicy POLICY ImageIdsBatchPolicy() def test_per_sample_independent(): bh torch.tensor([64, 64, 96, 64]); bw torch.tensor([64, 64, 64, 96]) out POLICY.prepare(bh, bw) assert out.shape[0] 4 # batch 维度对 # 第 3 个样本 (h96) 的 id 数应不同于第 1 个 (h64) assert out[2].shape[0] ! out[0].shape[0] # 独立计算不共享 def test_broadcast_bug_absent(): bh torch.tensor([64, 64]); bw torch.tensor([64, 64]) out POLICY.prepare(bh, bw) # 两个样本 id 应分别由各自高宽算不是共享第 0 个 assert torch.equal(out[0], out[1]) or out.shape[0] 2 def test_shape_matches_transformer_tokens(): # id 序列长度应等于 (h/patch)*(w/patch) bh torch.tensor([64]); bw torch.tensor([64]) out POLICY.prepare(bh, bw) assert out[0].shape[0] (64 // 2) * (64 // 2) def test_variable_size_rejected_when_disallowed(): bh torch.tensor([64, 96]); bw torch.tensor([64, 64]) problems POLICY.validate_batch(bh, bw, POLICY.prepare(bh, bw)) assert any(尺寸不一致 in p for p in problems) or POLICY.allow_variable_size def test_training_bs4_no_shape_error(): bh torch.tensor([64, 64, 64, 64]); bw torch.tensor([64, 64, 64, 64]) out POLICY.prepare(bh, bw) assert out.shape (4, (64 // 2) * (64 // 2), 2)CI 把test_per_sample_independent与test_training_bs4_no_shape_error作为 FLUX.2 img2img 训练脚本的必过项要求「batch1 时每个样本位置 id 独立、形状匹配」。八、排查清单_prepare_image_idsbatch 处理错按顺序查batch1 正常、batch1 异常重点查_prepare_image_ids是否在 batch 维度共享/广播了单样本 id。是否每个样本独立算 id正确做法是for i: ids[i]prepare(h[i],w[i])不是prepare(batch_h, batch_w)。报错expanded size ... must matchbatch 内某样本高宽与第 0 个差太多broadcast 出的 id 长度与 token 数不符。id 维度是否正确h_id, w_idFLUX.2 位置编码依赖二维 id错维度会让位置编码全乱。batch 内尺寸是否一致不一致需allow_variable_size pad否则 stack 失败。shape 是否与 transformer 期望 token 数一致(h/patch)*(w/patch)必须等于 latent token 数。九、小结「Fix incorrect batch handling in _prepare_image_ids usage」本质是图像位置 id 是「每样本独立」的量取决于各自高宽但训练脚本在 batch 维度上错误地把它当全局共享/批量广播导致 batch1 时位置编码错位、loss 异常或 shape 报错。第一层改为「每个样本独立计算 正确堆叠」第二层把 batch 处理契约每样本独立、id 维度、尺寸一致性收敛到ImageIdsBatchPolicy单一真源validate_batch校验第三层用 pytest 守住「batch 内每样本 id 独立、形状匹配、变量尺寸受控」。通用教训**任何「与单样本几何高宽/位置强相关」的张量位置 id、RoPE、mask在 batch 维度都必须逐样本独立计算再堆叠绝不能批量广播或共享否则 batch1 即隐性错位。