48 lines
1.2 KiB
Python
48 lines
1.2 KiB
Python
import torch
|
||
|
||
def get_device():
|
||
"""
|
||
获取可用的设备（GPU或CPU）
|
||
|
||
返回:
|
||
torch.device: 可用的设备
|
||
"""
|
||
if torch.cuda.is_available():
|
||
return torch.device('cuda')
|
||
else:
|
||
return torch.device('cpu')
|
||
|
||
def to_device(data, device):
|
||
"""
|
||
将数据移动到指定设备
|
||
|
||
参数:
|
||
data: 要移动的数据（可以是张量、列表、元组或字典）
|
||
device: 目标设备
|
||
|
||
返回:
|
||
移动到设备上的数据
|
||
"""
|
||
if isinstance(data, (list, tuple)):
|
||
return [to_device(x, device) for x in data]
|
||
elif isinstance(data, dict):
|
||
return {k: to_device(v, device) for k, v in data.items()}
|
||
elif isinstance(data, torch.Tensor):
|
||
return data.to(device)
|
||
else:
|
||
return data
|
||
|
||
class DeviceDataLoader:
|
||
"""
|
||
将DataLoader的数据移动到指定设备的包装器
|
||
"""
|
||
def __init__(self, dataloader, device):
|
||
self.dataloader = dataloader
|
||
self.device = device
|
||
|
||
def __iter__(self):
|
||
for batch in self.dataloader:
|
||
yield to_device(batch, self.device)
|
||
|
||
def __len__(self):
|
||
return len(self.dataloader) |