深度学习 GPU的调用

pytorch 的GPU 使用

train_on_gpu = torch.cuda.is_available()

if not train_on_gpu:
    print('CUDA is not available.  Training on CPU ...')
else:
    print('CUDA is available!  Training on GPU ...')
    
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

模型也得放到你的CPU或者GPU

model_ft =model.to(device)

tensorflow gpu 的使用

直接全局使用

gpus = tf.config.experimental.list_physical_devices(device_type='GPU')
for gpu in gpus:
	tf.config.experimental.set_memory_growth(gpu, True)