search

bigmb / Unet-Segmentation-Pytorch-Nest-of-Unets

Implementation of different kinds of Unet Models for Image Segmentation - Unet , RCNN-Unet, Attention Unet, RCNN-Attention Unet, Nested Unet
MIT License
1.87k stars 345 forks source link

Why can't train the model on the GPU #26

Closed JasonChenhx closed 4 years ago

JasonChenhx commented 4 years ago

I don't modify the code about GPU, and "CUDA is available. Training on GPU". What's more, I have put the input and model into the cuda. So What else do I need to put in cuda? Thanks.

bigmb commented 4 years ago

What does the output to 'device' give you? If it's cuda:0 then its a GPU, otherwise it's being trained on CPU.

Other than that, everything is there for training on a GPU or a CPU, according to the device found.

JasonChenhx commented 4 years ago

What does the output to 'device' give you? If it's cuda:0 then its a GPU, otherwise it's being trained on CPU.

Other than that, everything is there for training on a GPU or a CPU, according to the device found.

The run result is cuda:0, and CUDA is available. Training on GPU" And I have put the input、target、model into the cuda.

bigmb commented 4 years ago

They all will be trained in the GPU and the will be transferred back to the CPU. Look into the line 244 in pytorch_run It transfers your training data to GPU and then as you see in line 298 it will be called back to the CPU.