search

nv-nguyen / gigapose

[CVPR 2024] PyTorch implementation of GigaPose: Fast and Robust Novel Object Pose Estimation via One Correspondence
https://nv-nguyen.github.io/gigaPose/
MIT License
144 stars 12 forks source link

AENet and ISTNet training #7

Closed DanieleBertagnoli closed 8 months ago

DanieleBertagnoli commented 8 months ago

Hi, while I was looking to the code I noticed that the training of both AENet and ISTNet is done in gigapose.py. It couldn't be done directly in the ae_net.py and ist_net.py files? Or there is a major reason behind this choice? As you offer the possibility to train the two networks even separately, I guess that it is possible to separate the training code of the two nets from gigapose.py. That's just a curiosity since I was wondering if in my project (just to be coherent with my project's design) this split could be done without causing any misbehavior.

Best regards!

nv-nguyen commented 8 months ago

I put the training_step of both AENet and ISTNet in gigapose.py to avoid repetitions since they are supposed to use the same training data.

I add an option nets_to_train in the config file so that you can train them separately by simply a flag nets_to_train=ae or nets_to_train=ist such as python train.py nets_to_train=ae

DanieleBertagnoli commented 8 months ago

Ok, I'm just trying to figure out how the models are trained. If I'm not wrong, in gigapose.py the functions related to training are:

I would like to recreate an AENet class and ISTNet class such that they will be completely separated from gigapose.py (in this way those two models can be used as blackbox in gigapose), I think this could be usefull to use the model with custom objects.

nv-nguyen commented 8 months ago

Yes, they are correct.

As I mentioned above, if you run python train.py nets_to_train=ae, it will not affect ISTNet and same for running python train.py nets_to_train=ist will not affect AENet so I would say the networks have been already separated. You do not need to recreate two separate classes again.

DanieleBertagnoli commented 8 months ago

There is no explicit need of that, you're right but I would like to have the AENet checkpoint separated with respect to the ISTNet one. Maybe is already possible and I'm missing something?

nv-nguyen commented 8 months ago

Why AENet checkpoint should be separated to ISTNet one? You can simply load both and use only one of them, no?

nv-nguyen commented 8 months ago

Closed the issue due to inactivity. Feel free to re-open it!