uncbiag / uniGradICON

The official website for uniGradICON: A Foundation Model for Medical Image Registration
Apache License 2.0
64 stars 3 forks source link

The network weight file does not match the network. #14

Closed Promise13 closed 3 months ago

Promise13 commented 3 months ago

Hello, I'm very interested in your research content, but I've encountered a small issue. I obtained a network weight file named Step_2_final.trch from the online demo of unigradICON, and the link is as follows: Google Drive link. We found the training code for unigradICON at this GitHub repository. However, when we passed the network weight file Step_2_final.trch into the training code of unigradICON, we encountered a problem of mismatch between the network and the weights. What could be the reason for this?

lintian-a commented 3 months ago

Hi @Promise13 ,

Thanks for your interest in our work!

The weight we released is the complete model including the last step. Thus, you need to set include_last_step to True in make_network(input_shape, include_last_step=True, framework=framework) .

In the training code, we train the model make_network(input_shape, include_last_step=False, framework=framework) first (Line 272-288). Then we train the model with the last step (Line 290-316) with make_network(input_shape, include_last_step=True, framework=framework).

Promise13 commented 3 months ago

The training script ran successfully. Thank you for your response.

Promise13 commented 3 months ago

You mentioned that the training requires two steps, and we only have the network weights for the second step. How can we find the network weights for the first step?

lintian-a commented 3 months ago

The weight file we released includes the weight both for the first step (three UNets with multi-resolution scheme) and the second step (one UNet at the full resolution).

To better assist you, may I get some context here? For example, why do you only want the first step? Are you trying to finetune the model?

lintian-a commented 3 months ago

Hi @Promise13 ,

Do you still need help with training/finetuning uniGradICON?

Promise13 commented 3 months ago

Sorry, I didn't see your reply. I really appreciate your kind help, but I don't need assistance at the moment.