Python Contextの例

コード例 #1

ファイル: __init__.py プロジェクト: XJTUeducation/nnabla-ext-cuda

def context(device_id=0, type_config='float', **kw):
    """CUDA context."""
    backends = ['cuda:float', 'cpu:float']
    if type_config == 'half':
        backends = ['cuda:half', 'cuda:float', 'cpu:float']
    elif type_config == 'mixed_half':
        backends = ['cuda:mixed_half', 'cuda:float', 'cpu:float']
    elif type_config == 'float':
        pass
    else:
        raise ValueError("Unknown data type config is given %s" % type_config)
    return Context(backends, array_classes()[0], device_id=str(device_id))

コード例 #2

ファイル: __init__.py プロジェクト: zw0610/nnabla-ext-cuda

def context(device_id='0', type_config='float', *kw):
    """CUDNN context"""
    from nnabla_ext.cuda import array_classes
    backends = ['cudnn:float', 'cuda:float', 'cpu:float']
    if type_config == 'half':
        backends = [
            'cudnn:half', 'cudnn:float', 'cuda:half', 'cuda:float', 'cpu:float'
        ]
    elif type_config == 'mixed_half':
        backends = [
            'cudnn:mixed_half', 'cudnn:half', 'cudnn:float', 'cuda:mixed_half',
            'cuda:half', 'cuda:float', 'cpu:float'
        ]
    elif type_config == 'float':
        pass
    else:
        raise ValueError("Unknown data type config is given %s" % type_config)
    return Context(backends, array_classes()[0], device_id=str(device_id))

コード例 #3

ファイル: __init__.py プロジェクト: stevevista/nbla-go-zero

def context(device_id=0, *kw):
    """CUDNN context"""
    from nnabla_ext.cuda import array_classes
    return Context('cpu|cuda', array_classes()[0], device_id=str(device_id), compute_backend='default|cudnn')

コード例 #4

ファイル: __init__.py プロジェクト: stevevista/nbla-go-zero

def context(device_id=0, **kw):
    """CUDA context."""
    return Context('cpu|cuda',
                   array_classes()[0],
                   device_id=str(device_id),
                   compute_backend='default')