hli1221 / imagefusion-rfn-nest

RFN-Nest(Information Fusion, 2021, Highly Cited Paper) - PyTorch =1.5,Python=3.7
108 stars 18 forks source link

About the code: #4

Open softwarePupil opened 2 years ago

softwarePupil commented 2 years ago

Thank you for open sourcing your code! But here's a question in train_fusionnet.py, should we add this line of code when we start training in the second epoch to put our model in the GPU?

print('Epoch %d.....' % e)
image_set_ir, batches = utils.load_dataset(original_imgs_path, batch_size)
fusion_model.train()
fusion_model.cuda()
count = 0

Thank you for your attention!

Gao907 commented 2 years ago

Thank you for open sourcing your code! But here's a question in train_fusionnet.py, should we add this line of code when we start training in the second epoch to put our model in the GPU?

print('Epoch %d.....' % e)
image_set_ir, batches = utils.load_dataset(original_imgs_path, batch_size)
fusion_model.train()
fusion_model.cuda()
count = 0

Thank you for your attention!

I met the same problem as yours. I've tried to put the fusion_model on GPU and then the training script can run.

hli1221 commented 1 year ago

Hi, thanks for your attention. This also works for the program. When execute 'utils.load_dataset()', the program will randomly choose the training data from whole dataset.