Python torch.set_default_tensor_type方法代码示例

# 需要导入模块: import torch [as 别名]
# 或者: from torch import set_default_tensor_type [as 别名]
def _test():
    device = torch.device('cuda')
    torch.set_default_tensor_type('torch.cuda.FloatTensor')
    dataset = DiamondDataset(num_points=int(1e6), width=20, bound=2.5, std=0.04)

    from utils import torchutils
    from matplotlib import pyplot as plt
    data = torchutils.tensor2numpy(dataset.data)
    fig, ax = plt.subplots(1, 1, figsize=(5, 5))
    # ax.scatter(data[:, 0], data[:, 1], s=2, alpha=0.5)
    bound = 4
    bounds = [[-bound, bound], [-bound, bound]]
    # bounds = [
    #     [0, 1],
    #     [0, 1]
    # ]
    ax.hist2d(data[:, 0], data[:, 1], bins=256, range=bounds)
    ax.set_xlim(bounds[0])
    ax.set_ylim(bounds[1])
    plt.show() 

# 需要导入模块: import torch [as 别名]
# 或者: from torch import set_default_tensor_type [as 别名]
def set_device(use_gpu, multi_gpu, _log):
    # Decide which device to use.
    if use_gpu and not torch.cuda.is_available():
        raise RuntimeError('use_gpu is True but CUDA is not available')

    if use_gpu:
        device = torch.device('cuda')
        torch.set_default_tensor_type('torch.cuda.FloatTensor')
    else:
        device = torch.device('cpu')

    if multi_gpu and torch.cuda.device_count() == 1:
        raise RuntimeError('Multiple GPU training requested, but only one GPU is available.')

    if multi_gpu:
        _log.info('Using all {} GPUs available'.format(torch.cuda.device_count()))

    return device 

# 需要导入模块: import torch [as 别名]
# 或者: from torch import set_default_tensor_type [as 别名]
def __init__(self, **kwargs):
        super(PyTorchExecutor, self).__init__(**kwargs)

        self.global_training_timestep = 0

        self.cuda_enabled = torch.cuda.is_available()

        # In PyTorch, tensors are default created on the CPU unless assigned to a visible CUDA device,
        # e.g. via x = tensor([0, 0], device="cuda:0") for the first GPU.
        self.available_devices = os.environ.get("CUDA_VISIBLE_DEVICES")
        # TODO handle cuda tensors

        self.default_torch_tensor_type = self.execution_spec.get("dtype", "torch.FloatTensor")
        if self.default_torch_tensor_type is not None:
            torch.set_default_tensor_type(self.default_torch_tensor_type)

        self.torch_num_threads = self.execution_spec.get("torch_num_threads", 1)
        self.omp_num_threads = self.execution_spec.get("OMP_NUM_THREADS", 1)

        # Squeeze result dims, often necessary in tests.
        self.remove_batch_dims = True 

# 需要导入模块: import torch [as 别名]
# 或者: from torch import set_default_tensor_type [as 别名]
def initialize(self, fixed=None):

        # Parse options
        self.args = self.parse(fixed)

        # Setting default torch Tensor type
        if self.args.cuda and torch.cuda.is_available():
            torch.set_default_tensor_type('torch.cuda.FloatTensor')
            cudnn.benchmark = True
        else:
            torch.set_default_tensor_type('torch.FloatTensor')

        # Create weights saving directory
        if not os.path.exists(self.args.save_dir):
            os.mkdir(self.args.save_dir)

        # Create weights saving directory of target model
        model_save_path = os.path.join(self.args.save_dir, self.args.exp_name)

        if not os.path.exists(model_save_path):
            os.mkdir(model_save_path)

        return self.args 

# 需要导入模块: import torch [as 别名]
# 或者: from torch import set_default_tensor_type [as 别名]
def BindsNET_cpu(n_neurons, time):
    t0 = t()

    torch.set_default_tensor_type("torch.FloatTensor")

    t1 = t()

    network = Network()
    network.add_layer(Input(n=n_neurons), name="X")
    network.add_layer(LIFNodes(n=n_neurons), name="Y")
    network.add_connection(
        Connection(source=network.layers["X"], target=network.layers["Y"]),
        source="X",
        target="Y",
    )

    data = {"X": poisson(datum=torch.rand(n_neurons), time=time)}
    network.run(inputs=data, time=time)

    return t() - t0, t() - t1 

# 需要导入模块: import torch [as 别名]
# 或者: from torch import set_default_tensor_type [as 别名]
def BindsNET_gpu(n_neurons, time):
    if torch.cuda.is_available():
        t0 = t()

        torch.set_default_tensor_type("torch.cuda.FloatTensor")

        t1 = t()

        network = Network()
        network.add_layer(Input(n=n_neurons), name="X")
        network.add_layer(LIFNodes(n=n_neurons), name="Y")
        network.add_connection(
            Connection(source=network.layers["X"], target=network.layers["Y"]),
            source="X",
            target="Y",
        )

        data = {"X": poisson(datum=torch.rand(n_neurons), time=time)}
        network.run(inputs=data, time=time)

        return t() - t0, t() - t1 

# 需要导入模块: import torch [as 别名]
# 或者: from torch import set_default_tensor_type [as 别名]
def __init__(self, flags):

        torch.set_default_tensor_type('torch.cuda.FloatTensor')

        # fix the random seed or not
        fix_seed()

        self.setup_path(flags)

        self.network = mlp.MLPNet(num_classes=flags.num_classes)

        self.network = self.network.cuda()

        print(self.network)
        print('flags:', flags)

        if not os.path.exists(flags.logs):
            os.mkdir(flags.logs)

        flags_log = os.path.join(flags.logs, 'flags_log.txt')
        write_log(flags, flags_log)

        self.load_state_dict(flags.state_dict)

        self.configure(flags) 

# 需要导入模块: import torch [as 别名]
# 或者: from torch import set_default_tensor_type [as 别名]
def _prepare_device(self, n_gpu_use):
        """
        setup GPU device if available, move model into configured device
        """
        n_gpu = torch.cuda.device_count()
        if n_gpu_use > 0 and n_gpu == 0:
            self.logger.warning("Warning: There\'s no GPU available on this machine,"
                                "training will be performed on CPU.")
            n_gpu_use = 0
        if n_gpu_use > n_gpu:
            self.logger.warning("Warning: The number of GPU\'s configured to use is {}, but only {} are available "
                                "on this machine.".format(n_gpu_use, n_gpu))
            n_gpu_use = n_gpu
        device = torch.device('cuda:0' if n_gpu_use > 0 else 'cpu')
        if device.type == 'cuda':
            torch.set_default_tensor_type('torch.cuda.FloatTensor')
        list_ids = list(range(n_gpu_use))
        return device, list_ids 

# 需要导入模块: import torch [as 别名]
# 或者: from torch import set_default_tensor_type [as 别名]
def init_torch(rng_seed, cuda_seed):
    """
    Initializes the seeds for ALL potential randomness, including torch, numpy, and random packages.

    Args:
        rng_seed (int): the shared random seed to use for numpy and random
        cuda_seed (int): the random seed to use for pytorch's torch.cuda.manual_seed_all function
    """

    # default tensor
    torch.set_default_tensor_type('torch.cuda.FloatTensor')

    # seed everything
    torch.manual_seed(rng_seed)
    np.random.seed(rng_seed)
    random.seed(rng_seed)
    torch.cuda.manual_seed_all(cuda_seed)

    # make the code deterministic
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False 

# 需要导入模块: import torch [as 别名]
# 或者: from torch import set_default_tensor_type [as 别名]
def eye(m, dtype=None, device=None):
    """Returns a sparse matrix with ones on the diagonal and zeros elsewhere.

    Args:
        m (int): The first dimension of corresponding dense matrix.
        dtype (`torch.dtype`, optional): The desired data type of returned
            value vector. (default is set by `torch.set_default_tensor_type()`)
        device (`torch.device`, optional): The desired device of returned
            tensors. (default is set by `torch.set_default_tensor_type()`)

    :rtype: (:class:`LongTensor`, :class:`Tensor`)
    """

    row = torch.arange(m, dtype=torch.long, device=device)
    index = torch.stack([row, row], dim=0)

    value = torch.ones(m, dtype=dtype, device=device)

    return index, value 

# 需要导入模块: import torch [as 别名]
# 或者: from torch import set_default_tensor_type [as 别名]
def test_rnn_packed_sequence(self):
        num_layers = 2
        rnn = nn.RNN(input_size=self.h, hidden_size=self.h, num_layers=num_layers)
        for typ in [torch.float, torch.half]:
            x = torch.randn((self.t, self.b, self.h), dtype=typ).requires_grad_()
            lens = sorted([random.randint(self.t // 2, self.t) for _ in range(self.b)],
                          reverse=True)
            # `pack_padded_sequence` breaks if default tensor type is non-CPU
            torch.set_default_tensor_type(torch.FloatTensor)
            lens = torch.tensor(lens, dtype=torch.int64, device=torch.device('cpu'))
            packed_seq = nn.utils.rnn.pack_padded_sequence(x, lens)
            torch.set_default_tensor_type(torch.cuda.FloatTensor)
            hidden = torch.zeros((num_layers, self.b, self.h), dtype=typ)
            output, _ = rnn(packed_seq, hidden)
            self.assertEqual(output.data.type(), HALF)
            output.data.float().sum().backward()
            self.assertEqual(x.grad.dtype, x.dtype) 

# 需要导入模块: import torch [as 别名]
# 或者: from torch import set_default_tensor_type [as 别名]
def test_jv(alpha):

    torch.manual_seed(1)
    torch.set_default_tensor_type('torch.DoubleTensor')

    for _ in range(30):
        x = Variable(torch.randn(15))
        dout = torch.randn(15)

        y_hat = FusedProxFunction(alpha=alpha)(x).data


        ref = _fused_prox_jacobian(y_hat, dout)
        din_slow = fused_prox_jv_slow(y_hat, dout)
        din_fast = fused_prox_jv_fast(y_hat, dout)
        assert_allclose(ref.numpy(), din_slow.numpy(), atol=1e-5)
        assert_allclose(ref.numpy(), din_fast.numpy(), atol=1e-5) 

# 需要导入模块: import torch [as 别名]
# 或者: from torch import set_default_tensor_type [as 别名]
def nonlocalnet(input_layer,input_channel):
    if torch.cuda.is_available():
        torch.set_default_tensor_type('torch.cuda.FloatTensor')
        net = NONLocalBlock3D(in_channels=input_channel,mode='embedded_gaussian')
        out = net(input_layer)
    else:
        net = NONLocalBlock3D(in_channels=input_channel,mode='embedded_gaussian')
        out = net(input_layer)
    return out 

# 需要导入模块: import torch [as 别名]
# 或者: from torch import set_default_tensor_type [as 别名]
def get_model(args):
    torch.set_default_tensor_type('torch.FloatTensor')
    torch.manual_seed(args.seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(args.seed)

    model_class = {
        'aac': AdvantageActorCritic,
        'aac_lstm': AdvantageActorCriticLSTM,
        'aac_noisy': AdvantageActorCriticNoisy,
        'aac_depth': AdvantageActorCriticDepth,
        'aac_map': AdvantageActorCriticMap,
        'ppo': PPO,
        'ppo_map': PPOMap,
        'ppo_screen': PPOScreen,
        'mcts': MCTSPolicy,
        'state': StateBase,
        'es': ESBase,
        'planner': Planner
    }

    #
    # if model class derived from nn.Module then convert it to the current device
    # and load parameters if needed
    if issubclass(model_class[args.model], torch.nn.Module):
        model = Model.create(model_class[args.model], args, args.load)
    else:
        model = model_class[args.model](args)

    return model 

# 需要导入模块: import torch [as 别名]
# 或者: from torch import set_default_tensor_type [as 别名]
def setUp(self):
        torch.set_default_tensor_type(torch.DoubleTensor)
        torch.set_default_dtype(torch.float64)