ha-ha-ha-han / UKBiobank_deep_pretrain

Pretrained neural networks for UK Biobank brain MRI images. SFCN, 3D-ResNet etc.

MIT License

129 stars 36 forks source link

cpu version #2

Closed zhhsf closed 4 years ago

zhhsf commented 4 years ago

Hi, This toolbox is wonderful and really fast . Would you please provide a cpu version of it ?

ha-ha-ha-han commented 4 years ago

hi, you can remove all the model.cuda() lines, then the model will run with CPU automatically. E.g.

Example

model = SFCN() model = torch.nn.DataParallel(model) fp_ = './brain_age/run_20190719_00_epoch_best_mae.p' model.load_statedict(torch.load(fp)) ~~model.cuda()~~

zhhsf commented 4 years ago

Thanks for your reply. But the result of cpu vesion is not reproducible as cuda version.

model = SFCN() model = torch.nn.DataParallel(model) fp_ = './brain_age/run_20190719_00_epoch_bestmae.p' checkpoint = torch.load(fp,map_location='cpu') model.load_state_dict(checkpoint)

ha-ha-ha-han commented 4 years ago

Could you please give more details about the issue? Do you mean the CPU version does not perform well in UK Biobank data? Or there are some error messages?

I assume it is the 'CPU tensor broadcasting' bug? Fixed now ;) https://github.com/ha-ha-ha-han/UKBiobank_deep_pretrain/blob/master/examples.ipynb