tink/nerfstudio
4
0 Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-14 20:03:49 +08:00
Code Activity

bug fix for model loading from checkpoint (#1567)

Browse Source 
* bug fix for model loading from checkpoint

* removing prefix instead

* fix in load_pipeline if ddp as well

* fix

* formatting

---------

Co-authored-by: Jonáš Kulhánek <jonas.kulhanek@live.com>
This commit is contained in:
Terrance WangandJonáš Kulhánek authored and GitHub committed 2023-03-11 08:09:46 +01:00
1 parent 7ae2fce52b
commit 44ad94073d
1 file changed
+4 -4
Download Patch File Download Diff File Expand all files Collapse all files
nerfstudio/pipelines/base_pipeline.py
+4 -4
View File
@@ -107,9 +107,7 @@ class Pipeline(nn.Module):
return self.model.device
def load_state_dict(self, state_dict: Mapping[str, Any], strict: bool = True):
model_state = {
key.replace("_model.", ""): value for key, value in state_dict.items() if key.startswith("_model.")
}
model_state = {key[len("_model.") :]: value for key, value in state_dict.items() if key.startswith("_model.")}
pipeline_state = {key: value for key, value in state_dict.items() if not key.startswith("_model.")}
self._model.load_state_dict(model_state, strict=strict)
super().load_state_dict(pipeline_state, strict=False)
@@ -370,7 +368,9 @@ class VanillaPipeline(Pipeline):
loaded_state: pre-trained model state dict
step: training step of the loaded checkpoint
"""
state = {key.replace("module.", ""): value for key, value in loaded_state.items()}
state = {
(key[len("module.") :] if key.startswith("module.") else key): value for key, value in loaded_state.items()
}
self._model.update_to_step(step)
self.load_state_dict(state, strict=True)