Conditional Flow matching Cifar10

[nnna48]

I tried to change the unconditional flow matching algorithm for CIFAR10 to a conditional one, like with MNIST, but I failed and encountered unknown errors. How should I implement conditional CIFAR10? ...

[nnna48]

I met with this fault：
Windows fatal exception: access violation

Thread 0x0000754c (most recent call first):
File "D:\Anaconda\envs\FedDM\Lib\threading.py", line 359 in wait
File "D:\Anaconda\envs\FedDM\Lib\threading.py", line 655 in wait
File "D:\Anaconda\envs\FedDM\Lib\site-packages\tqdm_monitor.py", line 60 in run
File "D:\Anaconda\envs\FedDM\Lib\threading.py", line 1075 in _bootstrap_inner
File "D:\Anaconda\envs\FedDM\Lib\threading.py", line 1032 in _bootstrap

Thread 0x00009d6c (most recent call first):
File "C:\Users\24887\Desktop\Federate\conditional-flow-matching-main\torchcfm\models\unet\unet.py", line 327 in forward
File "D:\Anaconda\envs\FedDM\Lib\site-packages\torch\nn\modules\module.py", line 1747 in _call_impl
File "D:\Anaconda\envs\FedDM\Lib\site-packages\torch\nn\modules\module.py", line 1736 in _wrapped_call_impl
File "C:\Users\24887\Desktop\Federate\conditional-flow-matching-main\torchcfm\models\unet\unet.py", line 279 in _forward
File "C:\Users\24887\Desktop\Federate\conditional-flow-matching-main\torchcfm\models\unet\nn.py", line 122 in checkpoint
File "C:\Users\24887\Desktop\Federate\conditional-flow-matching-main\torchcfm\models\unet\unet.py", line 273 in forward
File "D:\Anaconda\envs\FedDM\Lib\site-packages\torch\nn\modules\module.py", line 1747 in _call_impl
File "D:\Anaconda\envs\FedDM\Lib\site-packages\torch\nn\modules\module.py", line 1736 in _wrapped_call_impl
File "C:\Users\24887\Desktop\Federate\conditional-flow-matching-main\torchcfm\models\unet\unet.py", line 69 in forward
File "D:\Anaconda\envs\FedDM\Lib\site-packages\torch\nn\modules\module.py", line 1747 in _call_impl
File "D:\Anaconda\envs\FedDM\Lib\site-packages\torch\nn\modules\module.py", line 1736 in _wrapped_call_impl
File "C:\Users\24887\Desktop\Federate\conditional-flow-matching-main\torchcfm\models\unet\unet.py", line 626 in forward
File "C:\Users\24887\Desktop\Federate\conditional-flow-matching-main\torchcfm\models\unet\unet.py", line 925 in forward
File "D:\Anaconda\envs\FedDM\Lib\site-packages\torch\nn\modules\module.py", line 1747 in _call_impl
File "D:\Anaconda\envs\FedDM\Lib\site-packages\torch\nn\modules\module.py", line 1736 in _wrapped_call_impl
File "C:\Users\24887\Desktop\Federate\conditional-flow-matching-main\examples\images\cifar10\utils_cifar.py", line 78 in
File "D:\Anaconda\envs\FedDM\Lib\site-packages\torchdiffeq_impl\misc.py", line 197 in forward
File "D:\Anaconda\envs\FedDM\Lib\site-packages\torch\nn\modules\module.py", line 1747 in _call_impl
File "D:\Anaconda\envs\FedDM\Lib\site-packages\torch\nn\modules\module.py", line 1736 in _wrapped_call_impl
File "D:\Anaconda\envs\FedDM\Lib\site-packages\torchdiffeq_impl\fixed_grid.py", line 10 in _step_func
File "D:\Anaconda\envs\FedDM\Lib\site-packages\torchdiffeq_impl\solvers.py", line 114 in integrate
File "D:\Anaconda\envs\FedDM\Lib\site-packages\torchdiffeq_impl\odeint.py", line 80 in odeint
File "C:\Users\24887\Desktop\Federate\conditional-flow-matching-main\examples\images\cifar10\utils_cifar.py", line 77 in generate_samples
File "C:\Users\24887\Desktop\Federate\conditional-flow-matching-main\examples\images\cifar10\train_cifar10.py", line 214 in train
File "D:\Anaconda\envs\FedDM\Lib\site-packages\absl\app.py", line 254 in _run_main
File "D:\Anaconda\envs\FedDM\Lib\site-packages\absl\app.py", line 308 in run
File "C:\Users\24887\Desktop\Federate\conditional-flow-matching-main\examples\images\cifar10\train_cifar10.py", line 230 in

进程已结束，退出代码为 -1073741819 (0xC0000005)

[nnna48]

the main change of my code：
def generate_samples(model, parallel, savedir, step, net_="normal", labels=None):
"""Save generated images (10 per class) for sanity check along training.

Parameters
----------
model: 神经网络模型
parallel: 是否并行训练
savedir: 保存路径
step: 训练步数
net_: 网络标识符
labels: 可选的条件标签Tensor
"""
model.eval()

model_ = copy.deepcopy(model)
if parallel:
    model_ = model_.module.to(device)
generated_class_list = torch.arange(10, device=device).repeat(3)
print(f"Batch size: {generated_class_list.shape}")
# 处理条件标签
if labels is None:
    # 默认生成10类，每类10张
    labels = torch.arange(10, device=device).repeat(10)
else:
    # 确保每个标签生成10张图像
    labels = torch.repeat_interleave(labels, 10)

num_samples = len(generated_class_list)
print(f"Number of samples: {num_samples}")
# num_samples = len(labels)

with torch.no_grad():
    # 使用torchdiffeq进行ODE求解
    traj = torchdiffeq.odeint(
        lambda t, x: model_(t, x, generated_class_list),
        torch.randn(num_samples, 3, 32, 32, device=device),
        torch.linspace(0, 1, 100, device=device),
        atol=1e-4,
        rtol=1e-4,
        method="euler",
    )
    traj = traj[-1].view([-1, 3, 32, 32]).clip(-1, 1)
    traj = traj / 2 + 0.5

save_image(traj, savedir + f"{net_}_generated_FM_images_step_{step}.png", nrow=10)
model.train()

net_model = UNetModelWrapper(
dim=(3, 32, 32),
num_res_blocks=2,
num_channels=FLAGS.num_channel,
channel_mult=[1, 2, 2, 2],
num_heads=4,
num_head_channels=64,
attention_resolutions="16",
dropout=0.1,
class_cond=True,
num_classes=10
).to(
device
) # new dropout + bs of 128