在PyTorch中,您可以使用.to(device)
方法将嵌套的字典中的所有支持的Tensor对象转移到GPU。以下是一个简单的例子
python
import torch
# 假设您已经有了一个名为device的GPU设备对象
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# 嵌套的字典,其中包含一些Tensors
nested_dict = {
'a': torch.randn(2, 2),
'b': {
'b1': torch.randn(2, 2),
'b2': torch.randn(2, 2)
},
'c': torch.randn(2, 2)
}
# 将嵌套字典中的所有Tensors移动到GPU
def to_gpu(data):
if isinstance(data, dict):
return {k: to_gpu(v) for k, v in data.items()}
elif isinstance(data, list):
return [to_gpu(i) for i in data]
elif isinstance(data, tuple):
return tuple([to_gpu(i) for i in data])
elif torch.is_tensor(data) and data.device != device:
return data.to(device)
else:
return data
nested_dict_gpu = to_gpu(nested_dict)
# 检查是否所有Tensors都已移动到GPU
for k, v in nested_dict_gpu.items():
if torch.is_tensor(v):
assert v.device == device
这个函数to_gpu
会递归地检查字典中的每个元素,如果是Tensor类型并且不在GPU上,就会使用.to(device)
方法转移它。您需要先设置device
变量指向您的GPU设备。如果没有GPU可用,它会默认使用CPU。