def set_caffe_mode(gpu):
''' Set whether caffe runs in gpu or not, input is boolean '''
if gpu:
caffe.set_mode_gpu()
else:
caffe.set_mode_cpu()
微信
公众号
扫码关注公众号