mirror of
https://github.com/SheffieldML/GPy.git
synced 2026-06-11 15:15:15 +02:00
adapt gpu initialization multiple gpu cards
This commit is contained in:
parent
cf33808673
commit
bdf478956d
2 changed files with 27 additions and 5 deletions
|
|
@ -6,6 +6,7 @@ import numpy as np
|
|||
from stationary import Stationary
|
||||
from psi_comp import PSICOMP_RBF
|
||||
from psi_comp.rbf_psi_gpucomp import PSICOMP_RBF_GPU
|
||||
from ...util.gpu_init import initGPU
|
||||
from ...util.config import *
|
||||
|
||||
class RBF(Stationary):
|
||||
|
|
@ -24,6 +25,7 @@ class RBF(Stationary):
|
|||
self.group_spike_prob = False
|
||||
self.psicomp = PSICOMP_RBF()
|
||||
if self.useGPU:
|
||||
initGPU()
|
||||
self.psicomp = PSICOMP_RBF_GPU()
|
||||
else:
|
||||
self.psicomp = PSICOMP_RBF()
|
||||
|
|
|
|||
|
|
@ -6,17 +6,37 @@ providing CUBLAS handle: cublas_handle
|
|||
"""
|
||||
|
||||
try:
|
||||
import pycuda.autoinit
|
||||
from scikits.cuda import cublas
|
||||
import scikits.cuda.linalg as culinalg
|
||||
culinalg.init()
|
||||
cublas_handle = cublas.cublasCreate()
|
||||
initSuccess = True
|
||||
except:
|
||||
initSuccess = False
|
||||
|
||||
gpu_initialized = False
|
||||
gpu_device = None
|
||||
gpu_context = None
|
||||
|
||||
def initGPU(gpuid=None):
|
||||
if gpu_initialized:
|
||||
return
|
||||
if gpuid==None:
|
||||
return pycuda.tools.make_default_context()
|
||||
try:
|
||||
import pycuda.autoinit
|
||||
gpu_initialized = True
|
||||
except:
|
||||
pass
|
||||
else:
|
||||
return pycuda.driver.Device(gpuid).make_context()
|
||||
try:
|
||||
import pycuda.driver
|
||||
pycuda.driver.init()
|
||||
if gpuid>=pycuda.driver.Device.count():
|
||||
return
|
||||
gpu_device = pycuda.driver.Device(gpuid)
|
||||
gpu_context = gpu_device.make_context()
|
||||
gpu_initialized = True
|
||||
except:
|
||||
pass
|
||||
|
||||
def closeGPU():
|
||||
if gpu_context is not None:
|
||||
gpu_context.detach()
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue