Python cuda.is_available方法代码示例

# 需要导入模块: from torch import cuda [as 别名]
# 或者: from torch.cuda import is_available [as 别名]
def __init__(self, gen_model,disc_model,use_cuda_if_available=True):
        super(BaseGanLearner,self).__init__()
        self.model_dir = os.getcwd()
        self.gen_model = gen_model
        self.disc_model = disc_model
        self.cuda = False
        if use_cuda_if_available and cuda.is_available():
            self.cuda = True

        self.__train_history__ = {}

        self.gen_optimizer = None
        self.disc_optimizer = None
        self.gen_running_loss = None
        self.disc_running_loss = None

        self.visdom_log = None
        self.tensorboard_log = None 

# 需要导入模块: from torch import cuda [as 别名]
# 或者: from torch.cuda import is_available [as 别名]
def __init__(self, dataset: Dataset, batch_size: int, source_names: List[str], target_names: List[str],
                 sort_key: Optional[Callable] = None, **kwargs):
        self.dataset = dataset
        self.source_names = source_names
        self.target_names = target_names
        # sort by the first field if no sort key is given
        if sort_key is None:
            def sort_key(x):
                return getattr(x, self.source_names[0])
        device = None if cuda.is_available() else -1
        self.dl = BucketIterator(dataset, batch_size=batch_size, sort_key=sort_key, device=device, **kwargs)
        self.bs = batch_size
        self.iter = 0 

# 需要导入模块: from torch import cuda [as 别名]
# 或者: from torch.cuda import is_available [as 别名]
def __init__(self, use_cuda_if_available=True):

        self.cuda = False
        self.fp16_mode = False
        if use_cuda_if_available and cuda.is_available():
            self.cuda = True
            cudnn.benchmark = True

        self.epoch_start_funcs = []
        self.batch_start_funcs = []
        self.epoch_end_funcs = []
        self.batch_end_funcs = []
        self.train_completed_funcs = [] 

# 需要导入模块: from torch import cuda [as 别名]
# 或者: from torch.cuda import is_available [as 别名]
def gpu_available() -> bool:
        return False 

# 需要导入模块: from torch import cuda [as 别名]
# 或者: from torch.cuda import is_available [as 别名]
def __init__(self, dataset, batch_size, sort_key, target_roles=None, max_context_size=130000, backwards=False,
                 **kwargs):
        self.target_roles = target_roles
        self.text_field = dataset.fields['text']
        self.max_context_size = max_context_size
        self.backwards = backwards
        device = None if cuda.is_available() else -1
        super().__init__(dataset=dataset, batch_size=batch_size, sort_key=sort_key, device=device, **kwargs) 

# 需要导入模块: from torch import cuda [as 别名]
# 或者: from torch.cuda import is_available [as 别名]
def check_cuda(torch_var, use_cuda=False):
    if use_cuda and cuda.is_available():
        return torch_var.cuda()
    else:
        return torch_var 

# 需要导入模块: from torch import cuda [as 别名]
# 或者: from torch.cuda import is_available [as 别名]
def get_info():
    """
    Get gpu info.

    :return: <dict> gpu info
    """
    return {
        "has_cuda": cuda.is_available(),
        "devices": [] if not cuda.is_available() else [cuda.get_device_name(i) for i in range(cuda.device_count())],
    } 

# 需要导入模块: from torch import cuda [as 别名]
# 或者: from torch.cuda import is_available [as 别名]
def generate_images(model, batch, mask_descriptors, num_samples=64, temp=1.,
                    verbose=False):
    """Generates image completions based on the images in batch masked by the
    masks in mask_descriptors. This will generate
    batch.size(0) * len(mask_descriptors) * num_samples completions, i.e.
    num_samples completions for every image and mask combination.

    Parameters
    ----------
    model : pixconcnn.models.pixel_constrained.PixelConstrained instance

    batch : torch.Tensor

    mask_descriptors : list of mask_descriptor
        See utils.masks.MaskGenerator for allowed mask_descriptors.

    num_samples : int
        Number of samples to generate for a given image-mask combination.

    temp : float
        Temperature for sampling.

    verbose : bool
        If True prints progress information while generating images
    """
    device = torch_device("cuda" if cuda_is_available() else "cpu")
    model.to(device)
    outputs = []
    for i in range(batch.size(0)):
        outputs_per_img = []
        for j in range(len(mask_descriptors)):
            if verbose:
                print("Generating samples for image {} using mask {}".format(i, mask_descriptors[j]))
            # Get image and mask combination
            img = batch[i:i+1]
            mask_generator = MaskGenerator(model.prior_net.img_size, mask_descriptors[j])
            mask = mask_generator.get_masks(1)
            # Create conditional pixels which will be used to sample completions
            cond_pixels = get_repeated_conditional_pixels(img, mask, model.prior_net.num_colors, num_samples)
            cond_pixels = cond_pixels.to(device)
            samples, log_probs = model.sample(cond_pixels, return_likelihood=True, temp=temp)
            outputs_per_img.append({
                "orig_img": img,
                "cond_pixels": cond_pixels,
                "mask": mask,
                "samples": samples,
                "log_probs": log_probs
            })
        outputs.append(outputs_per_img)
    return outputs