本文整理汇总了Python中torch.set_default_tensor_type方法的典型用法代码示例。如果您正苦于以下问题:Python torch.set_default_tensor_type方法的具体用法?Python torch.set_default_tensor_type怎么用?Python torch.set_default_tensor_type使用的例子?那么恭喜您, 这里精选的方法代码示例或许可以为您提供帮助。您也可以进一步了解该方法所在模块torch的用法示例。
在下文中一共展示了torch.set_default_tensor_type方法的23个代码示例,这些例子默认根据受欢迎程度排序。您可以为喜欢或者感觉有用的代码点赞,您的评价将有助于我们的系统推荐出更棒的Python代码示例。
示例1: _test
点赞 6
# 需要导入模块: 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()
开发者ID:bayesiains,项目名称:nsf,代码行数:22,
示例2: set_device
点赞 6
# 需要导入模块: 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
开发者ID:bayesiains,项目名称:nsf,代码行数:20,
示例3: __init__
点赞 6
# 需要导入模块: 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
开发者ID:rlgraph,项目名称:rlgraph,代码行数:23,
示例4: initialize
点赞 6
# 需要导入模块: 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
开发者ID:princewang1994,项目名称:TextSnake.pytorch,代码行数:25,
示例5: BindsNET_cpu
点赞 6
# 需要导入模块: 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(da