missing layers in the released checkpoint

[xiyichen]

Nice work! I'm trying to finetune the pretrained checkpoint on a custom dataset, but when I'm trying to load the pretrained weights from the released .safetensor into the model for finetuning, it seems a few dpt layers are missing in the pretrained model.

Do you know what caused this difference, or is there another model that I'm supposed to be loading for finetuning?

[xiyichen]

Setting strict to False seem to work temporarily, but I'm still curious if there's something I missed

[jedyang97]

Could you show the full steps you used to load these weights? I can try to see if I can reproduce the problem you are seeing.

[xiyichen]

Thanks for your reply. I only modified this part of the code in train.py:

fast3r/fast3r/train.py

Line 100 in 3ed2d93

if cfg.get("train"):

.

ckpt_path is the path to the .safetensors downloaded from huggingface. If I set strict=True, I get the error above.

[jedyang97]

Ah, I figured it out. In fast3r/croco/models/dpt_block.py, the make_scratch() function created a bunch of pointers:

fast3r/fast3r/croco/models/dpt_block.py

Lines 42 to 86 in b045ad6

	scratch.layer1_rn = nn.Conv2d(
	in_shape[0],
	out_shape1,
	kernel_size=3,
	stride=1,
	padding=1,
	bias=False,
	groups=groups,
	)
	scratch.layer2_rn = nn.Conv2d(
	in_shape[1],
	out_shape2,
	kernel_size=3,
	stride=1,
	padding=1,
	bias=False,
	groups=groups,
	)
	scratch.layer3_rn = nn.Conv2d(
	in_shape[2],
	out_shape3,
	kernel_size=3,
	stride=1,
	padding=1,
	bias=False,
	groups=groups,
	)
	scratch.layer4_rn = nn.Conv2d(
	in_shape[3],
	out_shape4,
	kernel_size=3,
	stride=1,
	padding=1,
	bias=False,
	groups=groups,
	)
	scratch.layer_rn = nn.ModuleList(
	[
	scratch.layer1_rn,
	scratch.layer2_rn,
	scratch.layer3_rn,
	scratch.layer4_rn,
	]
	)

such that,

model.net.downstream_head.dpt.scratch.layer_rn[0] points to model.net.downstream_head.dpt.scratch.layer1_rn
model.net.downstream_head.dpt.scratch.layer_rn[1] points to model.net.downstream_head.dpt.scratch.layer2_rn
...
model.net.downstream_head.dpt.scratch.layer_rn[3] points to model.net.downstream_head.dpt.scratch.layer4_rn

There is only one actual copy of the weight stored in model.safetensors, which is why the load_state_dict() function is complaining. And your solution to set strict=False is the correct solution - it is OK to ignore those missing keys in this case since the model.net.downstream_head.dpt.scratch.layer_rn[3] in the actual network are just pointers - as long as the actual weight tensor model.net.downstream_head.dpt.scratch.layer4_rn is loaded, we should be okay.

Here is the full code to make fine-tuning work with the HF model.safetensors, for future reference:

checkpoint_path = '/path/to/model.safetensors'
ckpt = load_file(checkpoint_path)
updated_ckpt = {f'net.{k}': v for k, v in ckpt.items()}
model.load_state_dict(updated_ckpt, strict=False)

To clarify, the original trainer.fit(model=model, datamodule=datamodule, ckpt_path=cfg.get("ckpt_path")) is intended to load from a pytorch-lightning-saved checkpoint. Thanks for testing it out with loading the HF safetensors checkpoint!

[xiyichen]

I see. Thanks a lot for the explanation!