train without projection layer

[immaturehearts]

Is there any way to train without the proposed projection layer? like the ablation study did?

[lukasHoel]

Yes, you can change the configuration in train.sh as follows:

--finetune-config.cross_frame_attention.unproj_reproj_mode "none"

[immaturehearts]

I did so, but got some error.
File "/data/ViewDiff/viewdiff/model/custom_attention_processor.py", line 29, in expand_batch return x.reshape(n // frames_per_batch, frames_per_batch, *other_dims) RuntimeError: shape '[0, 3, 16, 1280]' is invalid for input of size 20480

[lukasHoel]

Ah sorry, I also forgot about this detail. You also need to add these arguments:

--finetune-config.cross_frame_attention.n_cfa_down_blocks "0"
--finetune-config.cross_frame_attention.n_cfa_up_blocks "0"
--finetune-config.cross_frame_attention.no_cfa_in_mid_block

It will stil perform cross-frame-attention, because we set

--finetune-config.cross_frame_attention.mode "pretrained"

[immaturehearts]

Thank you so much!! It works! And btw, I wonder why the gpu memory usage is the same with or without the projection layer? It's not quite what I expected.

[lukasHoel]

The memory usage should not be the same. For me, train.sh utilizes around 42GB of memory with the projection layer, but 32GB of memory without the projection layer on one A100 GPU. I used the original train.sh script and modified it as discussed in this thread to get these numbers.

[immaturehearts]

Oh, I see. I used the train_small.sh and it's nearly the same. Thanks.

[lukasHoel]

That's expected, because train_small.sh is not using the full architecture. In particular, note the differences for these attributes:

--finetune-config.cross_frame_attention.num_3d_layers "5" (default) vs. "1" (small)
--finetune-config.cross_frame_attention.dim_3d_grid "128" (default) vs. "64" (small)
--finetune-config.cross_frame_attention.n_cfa_down_blocks "1" (default) vs. "0" (small)
--finetune-config.cross_frame_attention.n_cfa_up_blocks "1" (default) vs. "0" (small)

It means, that train_small.sh is already using only a very small 3D-grid in the bottleneck of the U-Net and thus does not use that much memory for this layer anyways. So completely removing it would not further reduce the memory.

[immaturehearts]

Okay, and also I wonder how many epochs did you train to get a good result.

[lukasHoel]

We trained for 60K iterations with 2xA100 GPUs using the train.sh script