Skip to content

Commit

Permalink
Update train.py
Browse files Browse the repository at this point in the history
  • Loading branch information
PeterL1n committed Nov 4, 2021
1 parent c22b354 commit 48effc9
Showing 1 changed file with 2 additions and 0 deletions.
2 changes: 2 additions & 0 deletions train.py
Original file line number Diff line number Diff line change
Expand Up @@ -329,6 +329,8 @@ def init_model(self):
{'params': self.model.backbone.parameters(), 'lr': self.args.learning_rate_backbone},
{'params': self.model.aspp.parameters(), 'lr': self.args.learning_rate_aspp},
{'params': self.model.decoder.parameters(), 'lr': self.args.learning_rate_decoder},
{'params': self.model.project_mat.parameters(), 'lr': self.args.learning_rate_decoder},
{'params': self.model.project_seg.parameters(), 'lr': self.args.learning_rate_decoder},
{'params': self.model.refiner.parameters(), 'lr': self.args.learning_rate_refiner},
])
self.scaler = GradScaler()
Expand Down

2 comments on commit 48effc9

@PeterL1n
Copy link
Owner Author

@PeterL1n PeterL1n commented on 48effc9 Nov 9, 2021

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This commit fixed a bug in train.py that caused training failure. Now we added model.project_mat and model.project_seg to the optimizer so that their parameters can be optimized.

Note that the training script used in our official training and paper didn't have this bug. It was only in the released script. So, our paper and pretrained models are not affected by this issue.

@Brothre23
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks so much for the bug fix!
I also encountered the problem where after about 8~9 epochs, the output of the network got all black, and the loss exploded.
Hope it does not happen again.

Please sign in to comment.