BiohazardVFX/4DGaussians
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #60 from hustvl/fix-ckpt

Browse Source 
fix-ckpt-loading
This commit is contained in:
Geralt_of_Rivia 2023-12-27 15:50:00 +08:00 committed by GitHub
commit bf1c002203
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 23 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
README.md
View File

@ -76,6 +76,18 @@ For training synthetic scenes such as `bouncingballs`, run
python train.py -s data/dnerf/bouncingballs --port 6017 --expname "dnerf/bouncingballs" --configs arguments/dnerf/bouncingballs.py
```
You can customize your training config through the config files.
# Checkpoint
Also, you can training your model with checkpoint.
```python
python train.py -s data/dnerf/bouncingballs --port 6017 --expname "dnerf/bouncingballs" --configs arguments/dnerf/bouncingballs.py --checkpoint_iterations 200 # change it.
```
Then load checkpoint with:
```python
python train.py -s data/dnerf/bouncingballs --port 6017 --expname "dnerf/bouncingballs" --configs arguments/dnerf/bouncingballs.py --start_checkpoint "output/dnerf/bouncingballs/chkpnt_coarse_200.pth"
# finestage: --start_checkpoint "output/dnerf/bouncingballs/chkpnt_fine_200.pth"
```
## Rendering
Run the following script to render the images.

2
submodules/depth-diff-gaussian-rasterization

@ -1 +1 @@
Subproject commit 2eb32ea251d3b339dab3af8b6fd78d7dec3caf8e
Subproject commit e49506654e8e11ed8a62d22bcb693e943fdecacf

14
train.py
View File

@ -45,9 +45,15 @@ def scene_reconstruction(dataset, opt, hyper, pipe, testing_iterations, saving_i
gaussians.training_setup(opt)
if checkpoint:
breakpoint()
(model_params, first_iter) = torch.load(checkpoint)
gaussians.restore(model_params, opt)
# breakpoint()
if stage == "coarse" and stage not in checkpoint:
print("start from fine stage, skip coarse stage.")
# process is in the coarse stage, but start from fine stage
return
if stage in checkpoint:
(model_params, first_iter) = torch.load(checkpoint)
gaussians.restore(model_params, opt)
bg_color = [1, 1, 1] if dataset.white_background else [0, 0, 0]
background = torch.tensor(bg_color, dtype=torch.float32, device="cuda")
@ -274,7 +280,7 @@ def scene_reconstruction(dataset, opt, hyper, pipe, testing_iterations, saving_i
if (iteration in checkpoint_iterations):
print("\n[ITER {}] Saving Checkpoint".format(iteration))
torch.save((gaussians.capture(), iteration), scene.model_path + "/chkpnt" + str(iteration) + ".pth")
torch.save((gaussians.capture(), iteration), scene.model_path + "/chkpnt" +f"_{stage}_" + str(iteration) + ".pth")
def training(dataset, hyper, opt, pipe, testing_iterations, saving_iterations, checkpoint_iterations, checkpoint, debug_from, expname):
# first_iter = 0
tb_writer = prepare_output_and_logger(expname)