Python pycuda.driver方法代码示例

# 需要导入模块: import pycuda [as 别名]
# 或者: from pycuda import driver [as 别名]
def test_module_export(self, tmpdir, input_example, module_name, df_type):
        """ Tests the module export.

            Args:
                tmpdir: Fixture which will provide a temporary directory.

                input_example: Input to be passed to TaylorNet.

                module_name: Name of the module (section in config file).

                df_type: Parameter denoting type of export to be tested.
        """
        # Create neural module instance.
        module = NeuralModule.import_from_config("tests/configs/test_deploy_export.yaml", module_name)
        # Generate filename in the temporary directory.
        tmp_file_name = str(tmpdir.mkdir("export").join(module_name))
        input_example = input_example.cuda() if input_example is not None else input_example
        # Test export.
        self.__test_export_route(
            module=module, out_name=tmp_file_name, mode=df_type, input_example=input_example,
        ) 

# 需要导入模块: import pycuda [as 别名]
# 或者: from pycuda import driver [as 别名]
def test_hf_bert(self, tmpdir, df_type):
        """ Tests BERT export.

            Args:
                tmpdir: Fixture which will provide a temporary directory.

                df_type: Parameter denoting type of export to be tested.
        """
        bert = nemo.collections.nlp.nm.trainables.common.huggingface.BERT(pretrained_model_name="bert-base-uncased")
        input_example = OrderedDict(
            [
                ("input_ids", torch.randint(low=0, high=16, size=(2, 16)).cuda()),
                ("token_type_ids", torch.randint(low=0, high=2, size=(2, 16)).cuda()),
                ("attention_mask", torch.randint(low=0, high=2, size=(2, 16)).cuda()),
            ]
        )
        # Generate filename in the temporary directory.
        tmp_file_name = str(tmpdir.mkdir("export").join("bert"))
        # Test export.
        self.__test_export_route(module=bert, out_name=tmp_file_name, mode=df_type, input_example=input_example) 

# 需要导入模块: import pycuda [as 别名]
# 或者: from pycuda import driver [as 别名]
def test_waveglow(self, tmpdir, df_type):
        url = "https://api.ngc.nvidia.com/v2/models/nvidia/waveglow_ljspeech/versions/2/files/WaveGlowNM.pt"
        ptfile = "./WaveGlowNM.pt"
        if not Path(ptfile).is_file():
            urllib.request.urlretrieve(url, ptfile)

        module = nemo_tts.WaveGlowInferNM(sample_rate=22050)
        module.restore_from(ptfile)
        module.eval()

        torch.manual_seed(1)
        mel = torch.randn(1, 80, 96).cuda()

        input_example = OrderedDict([("mel_spectrogram", mel)])
        tmp_file_name = str(tmpdir.mkdir("export").join("waveglow"))

        self.__test_export_route(module=module, out_name=tmp_file_name, mode=df_type, input_example=input_example) 

# 需要导入模块: import pycuda [as 别名]
# 或者: from pycuda import driver [as 别名]
def get_device_count(verbose=False):
    """
    Query device count through PyCuda.

    Arguments:
        verbose (bool): prints verbose logging if True, default False.

    Returns:
        int: Number of GPUs available.
    """
    try:
        import pycuda
        import pycuda.driver as drv
    except ImportError:
        if verbose:
            neon_logger.display("PyCUDA module not found")
        return 0
    try:
        drv.init()
    except pycuda._driver.RuntimeError as e:
        neon_logger.display("PyCUDA Runtime error: {0}".format(str(e)))
        return 0

    count = drv.Device.count()

    if verbose:
        neon_logger.display("Found {} GPU(s)".format(count))

    return count 

# 需要导入模块: import pycuda [as 别名]
# 或者: from pycuda import driver [as 别名]
def add_module(self, module):
            try:
                # LLVM can't produce CUBIN for some reason
                ptx = self._target_machine.emit_assembly(module)
                mod = pycuda.compiler.DynamicModule()
                mod.add_data(self._generated_builtins, pycuda.driver.jit_input_type.CUBIN, "builtins.cubin")
                mod.add_data(ptx.encode(), pycuda.driver.jit_input_type.PTX, module.name + ".ptx")
                ptx_mod = mod.link()

            except Exception as e:
                print("FAILED to generate PTX module:", e)
                print(ptx)
                return None

            self._modules[module] = ptx_mod 

# 需要导入模块: import pycuda [as 别名]
# 或者: from pycuda import driver [as 别名]
def get_compute_capability(device_id=None, verbose=False):
    """
    Query compute capability through PyCuda and check it's 5.0 (Maxwell) or
    greater.
    5.0 (GTX750 Ti) only fp32 support
    5.2 (GTX9xx series) required for fp16
    By default, check all devices and return the highest compute capability.

    Arguments:
        device_id (int): CUDA device id. Default to None, will iterate over
                         all devices if None.
        verbose (bool): prints verbose logging if True, default False.

    Returns:
        float: Zero if no GPU is found, otherwise highest compute capability.
    """
    try:
        import pycuda
        import pycuda.driver as drv
    except ImportError:
        if verbose:
            neon_logger.display("PyCUDA module not found")
        return 0
    try:
        drv.init()
    except pycuda._driver.RuntimeError as e:
        neon_logger.display("PyCUDA Runtime error: {0}".format(str(e)))
        return 0

    major_string = pycuda._driver.device_attribute.COMPUTE_CAPABILITY_MAJOR
    minor_string = pycuda._driver.device_attribute.COMPUTE_CAPABILITY_MINOR
    full_version = []
    if device_id is None:
        device_id = list(range(drv.Device.count()))
    elif isinstance(device_id, int):
        device_id = [device_id]

    for i in device_id:
        major = drv.Device(i).get_attribute(major_string)
        minor = drv.Device(i).get_attribute(minor_string)
        full_version += [major + minor / 10.]

    if verbose:
        neon_logger.display("Found GPU(s) with compute capability: {}".format(full_version))

    return max(full_version)