在训练sentinel这个数据集的时候，batch size只能设置为1

[kkc-nns]

作者您好，非常感谢您能够公开模型，这对我非常有帮助。
我在SEN12MS-CR数据集上面训练模型的时候，发现如果把sentinel.yaml中的batch_size直接改为2时会遇到下面的错误：
File "/mnt/workspace/users/codes/EMRDM-main/sgm/modules/diffusionmodules/sampling.py", line 465, in sampler_step x = x + ((1 - st_hat) / st_hat - (1 - st) / st) * mu + eps * append_dims(sigma_hat**2 - sigma**2, x.ndim) ** 0.5 RuntimeError: The size of tensor a (2) must match the size of tensor b (256) at non-singleton dimension 3
请问这个代码支持batch size > 1进行训练吗？我看到trainer的max_epochs是100，如果以batch size=1，在一台8卡H200上面训练，需要20天，而显存只占了10G/143G。或者是否需要训练完这100个epochs呢？
还想问一下，CUHK, CUHKv2, SEN2-MTC-new, SEN12MS-CR这四个数据集真实的训练epochs大概是多少呢？在yaml里面，他们的max_epochs分别是500, 2000, 1000, 100，是否都需要训练满这么多epochs?
非常感谢能得到您的解答！谢谢您！

[Ly403]

您好，首先非常感谢您对我们的工作感兴趣。

关于您的问题，首先代码是支持batch size > 1的训练的，你说的这个问题我之前发现过后来忘了修改。这个问题只在batch size > 1且使用随机采样的情况下会被触发，原因应该是这里面的系数（例如st, sigma等）没有广播到x的形状。很抱歉给您带来了困扰，也感谢您报告里问题，这两天我会赶紧修一下。您如果使用确定性采样是一定没有这个问题的，直接将s_churn, s_tmin, s_tmax, s_noise等注释了即可，如配置中这一行所示。您可以先训练时用确定性采样，测试时再尝试随机性采样。关于各个数据集训练的epoch，您可以在paper的附录里面的Figure 5查询到，如下：

希望这些能帮助到您！

[kkc-nns]

非常感谢您的回复，问题确实是st和sigma没有广播到x的形状。我尝试了确定性采样，的确可以规避这个问题。再次感谢您！