the total time of training the whole process

[aiXia121]

hi, thanks for your great work and sharing the code.
I'm really insterested in this work, and will do some tests and re-train your work.
i found generate_flows_v2.py to get all the flow's model costing a lot of time.
would you mind sharing how much time that you trained the whole process.
thanks

[aiXia121]

I am currently using PyTorch to work on accelerating the training process based on the code you provided, as the training with the original code is extremely time-consuming.

[maria-larchenko]

generate_flows_v2.py to get all the flow's model costing a lot of time

Hi! I think we spent around a week to generate flows on two RTX4090 and one RTX4060TI. Since then I found that partly this inefficiency comes from line where we call randpermute to sample from base_density and target_density:

modflows/src/neural_ode.py

Line 102 in 88de503

indx = torch.randperm(data_size)[:sample_size]

In other words, reducing INPUT_SIZE in generate_flows_v2.py notably speeds up a single flow. In the current version of code we import this variable from the encoder class, but in general it can be different from the encoder's one.

I measured time to generate a single flow with different base_density and target_density shapes. For RTX 4090 it scales as follows

INPUT_SIZE = 512
sample_size = 512 * 100
samples = 50_000
...
model = NeuralODE(input_dim=3, device=device, hidden=1024)  # 8195 params
...
>>> base_density    torch.Size([262144, 3])
>>> train_ode       05:31

INPUT_SIZE = 256
sample_size = 512 * 100
samples = 50_000
...
model = NeuralODE(input_dim=3, device=device, hidden=1024)
...
>>> base_density    torch.Size([65536, 3])
>>> train_ode       02:56

INPUT_SIZE = 256
sample_size = 512 * 100
samples = 50_000
...
model = NeuralODE(input_dim=3, device=device, hidden=64)  # 515 params
...
>>> base_density    torch.Size([65536, 3])
>>> train_ode       01:38

To generate 5826 rectified flows therefore we need around 05:31 * 2 * 5826 ~ 1100h ~ 44d. But we did it parallel, so it took less time. We also rectify the flow after the initial training, but you can try to spend less iterations on rectification step or even fully skip it and train the encoder on base_density and target_density_2 pairs. Actually, there is no need to save flow models at all. This modification also will speed up the encoder training. Then the time to reproduce the dataset for encoder B0 is 01:38 * 5826 ~ 158h ~ 6d if you are runnin generation in the single thread. B0 training itself takes ~1-2 days.