Function device_placement

modelscope/utils/device.py:40–82

Device placement function, allow user to specify which device to place model or tensor Args: framework (str): tensorflow or pytorch. device (str): gpu or cpu to use, if you want to specify certain gpu, use gpu:$gpu_id or cuda:$gpu_id. Returns: Context

(framework, device_name='gpu:0')

@contextmanager
def device_placement(framework, device_name='gpu:0'):
    """ Device placement function, allow user to specify which device to place model or tensor
    Args:
        framework (str):  tensorflow or pytorch.
        device (str):  gpu or cpu to use, if you want to specify certain gpu,
            use gpu:$gpu_id or cuda:$gpu_id.

    Returns:
        Context manager

    Examples:

        >>> # Requests for using model on cuda:0 for gpu
        >>> with device_placement('pytorch', device='gpu:0'):
        >>>     model = Model.from_pretrained(...)
    """
    device_type, device_id = verify_device(device_name)

    if framework == Frameworks.tf:
        import tensorflow as tf
        if device_type == Devices.gpu and not tf.test.is_gpu_available():
            logger.debug(
                'tensorflow: cuda is not available, using cpu instead.')
        device_type = Devices.cpu
        if device_type == Devices.cpu:
            with tf.device('/CPU:0'):
                yield
        else:
            if device_type == Devices.gpu:
                with tf.device(f'/device:gpu:{device_id}'):
                    yield

    elif framework == Frameworks.torch:
        import torch
        if device_type == Devices.gpu:
            if torch.cuda.is_available():
                torch.cuda.set_device(f'cuda:{device_id}')
            else:
                logger.debug(
                    'pytorch: cuda is not available, using cpu instead.')
        yield
    else:
        yield


def create_device(device_name):