Python cuda.get_device方法代码示例

# 需要导入模块: from chainer import cuda [as 别名]
# 或者: from chainer.cuda import get_device [as 别名]
def give_conditionalized_cell(self, src_batch, src_mask, noise_on_prev_word=False,
                                  demux=False):

        if self.lexical_probability_dictionary is not None:
            lexicon_probability_matrix = compute_lexicon_matrix(
                src_batch, self.lexical_probability_dictionary, self.Vo)
            if self.xp != np:
                lexicon_probability_matrix = cuda.to_gpu(
                    lexicon_probability_matrix, cuda.get_device(
                        self.dec.lin_o.W.data))
        else:
            lexicon_probability_matrix = None

        fb_concat = self.enc(src_batch, src_mask)

        mb_size, nb_elems, Hi = fb_concat.data.shape

        return self.dec.give_conditionalized_cell(fb_concat, src_mask,
                                                  noise_on_prev_word=noise_on_prev_word,
                                                  lexicon_probability_matrix=lexicon_probability_matrix,
                                                  lex_epsilon=self.lex_epsilon, demux=demux) 

# 需要导入模块: from chainer import cuda [as 别名]
# 或者: from chainer.cuda import get_device [as 别名]
def __init__(self, *, iterator, noise_iterator, optimizer_generator,
                 optimizer_critic, device=-1):

        if optimizer_generator.target.name is None:
            optimizer_generator.target.name = 'generator'

        if optimizer_critic.target.name is None:
            optimizer_critic.target.name = 'critic'

        iterators = {'main': iterator, 'z': noise_iterator}
        optimizers = {'generator': optimizer_generator,
                      'critic': optimizer_critic}

        super().__init__(iterators, optimizers, device=device)

        if device >= 0:
            cuda.get_device(device).use()
            [optimizer.target.to_gpu() for optimizer in optimizers.values()]

        self.xp = cuda.cupy if device >= 0 else np 

# 需要导入模块: from chainer import cuda [as 别名]
# 或者: from chainer.cuda import get_device [as 别名]
def __call__(self, x):
        """Applies the graph convolutional layer.

        Args:
            x: (~chainer.Variable): Input graph signal.

        Returns:
            ~chainer.Variable: Output of the graph convolution.
        """
        if self.has_uninitialized_params:
            with cuda.get_device(self._device_id):
                self._initialize_params(x.shape[1])
        if self.b is None:
            return self.func(x, self.W)
        else:
            return self.func(x, self.W, self.b) 

# 需要导入模块: from chainer import cuda [as 别名]
# 或者: from chainer.cuda import get_device [as 别名]
def forward_gpu(self, inputs):
        x, W = inputs[:2]
        n_batch, c_in, N = x.shape
        b = inputs[2] if len(inputs) == 3 else None
        xp = cuda.get_array_module(x)
        with cuda.get_device(x.data):
            K = self.K
            LmI_data, LmI_indices, LmI_indptr = self.LmI_tuple

            if x.dtype != LmI_data.dtype:
                LmI_data = LmI_data.astype(x.dtype)

            C = xp.empty((K, N, c_in, n_batch), dtype=x.dtype)
            chebyshev_matvec_gpu(C, x, K, n_batch,
                                 LmI_data, LmI_indices, LmI_indptr)

            C = C.transpose((3, 2, 0, 1))
            self.C = C
            y = xp.tensordot(C, W, ((1, 2), (1, 2)))

            if b is not None:
                y += b

            return xp.rollaxis(y, 2, 1),  # y.shape = (n_batch, c_out, N) 

# 需要导入模块: from chainer import cuda [as 别名]
# 或者: from chainer.cuda import get_device [as 别名]
def __call__(self, opt):
        if cuda.available:
            kernel = cuda.elementwise(
                'T p, T decay', 'T g', 'g += decay * p', 'weight_decay')

        rate = self.rate
        for name, param in opt.target.namedparams():
            if name == 'b' or name.endswith('/b'):
                continue
            p, g = param.data, param.grad
            with cuda.get_device(p) as dev:
                if int(dev) == -1:
                    g += rate * p
                else:
                    kernel(p, rate, g) 

# 需要导入模块: from chainer import cuda [as 别名]
# 或者: from chainer.cuda import get_device [as 别名]
def compute_lexicon_probability_matrix(self, src_batch):
        if self.lexical_probability_dictionary is not None:
            lexicon_probability_matrix = compute_lexicon_matrix(
                src_batch, self.lexical_probability_dictionary, self.Vo)
            if self.xp != np:
                lexicon_probability_matrix = cuda.to_gpu(
                    lexicon_probability_matrix, cuda.get_device(
                        self.dec.lin_o.W.data))
        else:
            lexicon_probability_matrix = None
        return lexicon_probability_matrix 

# 需要导入模块: from chainer import cuda [as 别名]
# 或者: from chainer.cuda import get_device [as 别名]
def init_state(self, param, state):
        xp = cuda.get_array_module(param.data)
        with cuda.get_device(param.data):
            state['m'] = xp.zeros_like(param.data)
            state['v'] = xp.zeros_like(param.data) 

# 需要导入模块: from chainer import cuda [as 别名]
# 或者: from chainer.cuda import get_device [as 别名]
def main():
    args = parse_args()

    _, log_file_exist = initialize_logging(
        logging_dir_path=args.save_dir,
        logging_file_name=args.logging_file_name,
        script_args=args,
        log_packages=args.log_packages,
        log_pip_packages=args.log_pip_packages)

    global_config.train = False

    num_gpus = args.num_gpus
    if num_gpus > 0:
        cuda.get_device(0).use()

    net = prepare_model(
        model_name=args.model,
        use_pretrained=args.use_pretrained,
        pretrained_model_file_path=args.resume.strip(),
        use_gpus=(num_gpus > 0))

    val_iterator, val_dataset_len = get_val_data_iterator(
        dataset_name=args.dataset,
        batch_size=args.batch_size,
        num_workers=args.num_workers)

    assert (args.use_pretrained or args.resume.strip())
    test(
        net=net,
        val_iterator=val_iterator,
        val_dataset_len=val_dataset_len,
        num_gpus=num_gpus,
        calc_weight_count=True,
        extended_log=True) 

# 需要导入模块: from chainer import cuda [as 别名]
# 或者: from chainer.cuda import get_device [as 别名]
def main():
    args = parse_args()

    _, log_file_exist = initialize_logging(
        logging_dir_path=args.save_dir,
        logging_file_name=args.logging_file_name,
        script_args=args,
        log_packages=args.log_packages,
        log_pip_packages=args.log_pip_packages)

    global_config.train = False

    num_gpus = args.num_gpus
    if num_gpus > 0:
        cuda.get_device(0).use()

    net = prepare_model(
        model_name=args.model,
        use_pretrained=args.use_pretrained,
        pretrained_model_file_path=args.resume.strip(),
        net_extra_kwargs={"aux": False, "fixed_size": False},
        use_gpus=(num_gpus > 0))

    test_dataset = get_test_dataset(
        dataset_name=args.dataset,
        dataset_dir=args.data_dir)

    assert (args.use_pretrained or args.resume.strip())
    test(
        net=net,
        test_dataset=test_dataset,
        num_gpus=num_gpus,
        num_classes=args.num_classes,
        calc_weight_count=True,
        extended_log=True,
        dataset_metainfo=get_metainfo(args.dataset)) 

# 需要导入模块: from chainer import cuda [as 别名]
# 或者: from chainer.cuda import get_device [as 别名]
def main():
    args = parse_args()

    _, log_file_exist = initialize_logging(
        logging_dir_path=args.save_dir,
        logging_file_name=args.logging_file_name,
        script_args=args,
        log_packages=args.log_packages,
        log_pip_packages=args.log_pip_packages)

    global_config.train = False

    num_gpus = args.num_gpus
    if num_gpus > 0:
        cuda.get_device(0).use()

    net = prepare_model(
        model_name=args.model,
        use_pretrained=args.use_pretrained,
        pretrained_model_file_path=args.resume.strip(),
        use_gpus=(num_gpus > 0))
    num_classes = net.classes if hasattr(net, "classes") else 1000
    input_image_size = net.in_size[0] if hasattr(net, "in_size") else args.input_size

    val_iterator, val_dataset_len = get_val_data_iterator(
        data_dir=args.data_dir,
        batch_size=args.batch_size,
        num_workers=args.num_workers,
        num_classes=num_classes)

    assert (args.use_pretrained or args.resume.strip())
    test(
        net=net,
        val_iterator=val_iterator,
        val_dataset_len=val_dataset_len,
        num_gpus=num_gpus,
        input_image_size=input_image_size,
        resize_inv_factor=args.resize_inv_factor,
        calc_weight_count=True,
        extended_log=True) 

# 需要导入模块: from chainer import cuda [as 别名]
# 或者: from chainer.cuda import get_device [as 别名]
def prepare_ch_context(num_gpus):
    use_gpus = (num_gpus > 0)
    if use_gpus:
        cuda.get_device(0).use()
    return use_gpus 

# 需要导入模块: from chainer import cuda [as 别名]
# 或者: from chainer.cuda import get_device [as 别名]
def set_net(self, net):
        self.source_net = deepcopy(net)
        self.target_net = deepcopy(net)
        if self.gpu:
            cuda.get_device(0).use()
            self.source_net.to_gpu()
            self.target_net.to_gpu()
        self.optimizer.setup(self.source_net)
        self.target_net.train = False 

# 需要导入模块: from chainer import cuda [as 别名]
# 或者: from chainer.cuda import get_device [as 别名]
def init_state(self, param):
        with cuda.get_device(param.data):
            self.state['s'] = [] 

# 需要导入模块: from chainer import cuda [as 别名]
# 或者: from chainer.cuda import get_device [as 别名]
def __call__(self, opt):
        if cuda.available:
            kernel = cuda.elementwise(
                'T low, T high', 
                'T p', 
                'p = (p < low) ? low : (p > high) ? high : p',
                'weight_clip')

        for param in opt.target.params():
            p = param.data
            with cuda.get_device(p) as dev:
                if int(dev) == -1:
                    numpy.clip(p, self.low, self.high)
                else:
                    kernel(self.low, self.high, p) 

# 需要导入模块: from chainer import cuda [as 别名]
# 或者: from chainer.cuda import get_device [as 别名]
def init_state(self, param, state):
        xp = cuda.get_array_module(param.data)
        with cuda.get_device(param.data):
            state['v'] = xp.zeros_like(param.data)