Don't convert namedtuple to tuple (#1589)

* Don't convert namedtuple to tuple

* Test namedtuples sent to device correctly
This commit is contained in:
Nathan Breitsch
2020-04-30 08:04:50 -04:00
committed by GitHub
parent d40425d257
commit 3eac6cfd4f
2 changed files with 17 additions and 4 deletions
+9 -4
View File
@@ -461,10 +461,15 @@ class TrainerDPMixin(ABC):
# when tuple
if isinstance(batch, tuple):
batch = list(batch)
for i, x in enumerate(batch):
batch[i] = self.__transfer_data_to_device(x, device, gpu_id)
return tuple(batch)
# when namedtuple
if hasattr(batch, '_fields'):
elem_type = type(batch)
return elem_type(*(self.__transfer_data_to_device(x, device, gpu_id) for x in batch))
else:
batch = list(batch)
for i, x in enumerate(batch):
batch[i] = self.__transfer_data_to_device(x, device, gpu_id)
return tuple(batch)
# when dict
if isinstance(batch, dict):
+8
View File
@@ -1,3 +1,4 @@
from collections import namedtuple
import platform
import pytest
@@ -221,6 +222,13 @@ def test_single_gpu_batch_parse():
assert batch[1][0]['b'].device.index == 0
assert batch[1][0]['b'].type() == 'torch.cuda.FloatTensor'
# namedtuple of tensor
BatchType = namedtuple('BatchType', ['a', 'b'])
batch = [BatchType(a=torch.rand(2, 3), b=torch.rand(2, 3)) for _ in range(2)]
batch = trainer.transfer_batch_to_gpu(batch, 0)
assert batch[0].a.device.index == 0
assert batch[0].a.type() == 'torch.cuda.FloatTensor'
def test_simple_cpu(tmpdir):
"""Verify continue training session on CPU."""