def _get_torchtext_data_iterator(include_lengths=False):
text_field = Field(
sequential=True,
pad_first=False, # nosec
init_token="<s>",
eos_token="</s>", # nosec
include_lengths=include_lengths
) # nosec
example1 = Example.fromdict({"text": "a b c a c"}, {"text": ("text", text_field)})
example2 = Example.fromdict({"text": "b c a a"}, {"text": ("text", text_field)})
example3 = Example.fromdict({"text": "c b a"}, {"text": ("text", text_field)})
dataset = Dataset(
[example1, example2, example3],
{"text": text_field},
)
text_field.build_vocab(dataset)
iterator = Iterator(
dataset,
batch_size=3,
sort_key=None,
device=None,
batch_size_fn=None,
train=True,
repeat=False,
shuffle=None,
sort=None,
sort_within_batch=None
)
return iterator, text_field
def get_dummy_torchtext_data_iterator(num_samples: int,
batch_size: int,
include_lengths: bool = False):
text_field = Field(
sequential=True,
pad_first=False, # nosec
init_token="<s>",
eos_token="</s>", # nosec
include_lengths=include_lengths,
) # nosec
dataset = Dataset(
[
Example.fromdict({"text": _generate_random_string()},
{"text": ("text", text_field)})
for _ in range(num_samples)
],
{"text": text_field},
)
text_field.build_vocab(dataset)
iterator = Iterator(
dataset,
batch_size=batch_size,
sort_key=None,
device=None,
batch_size_fn=None,
train=True,
repeat=False,
shuffle=None,
sort=None,
sort_within_batch=None,
)
return iterator, text_field
def test_single_gpu_batch_parse():
trainer = Trainer(gpus=1)
# non-transferrable types
primitive_objects = [
None, {}, [], 1.0, "x", [None, 2], {
"x": (1, 2),
"y": None
}
]
for batch in primitive_objects:
data = trainer.accelerator.batch_to_device(batch,
torch.device('cuda:0'))
assert data == batch
# batch is just a tensor
batch = torch.rand(2, 3)
batch = trainer.accelerator.batch_to_device(batch, torch.device('cuda:0'))
assert batch.device.index == 0 and batch.type() == 'torch.cuda.FloatTensor'
# tensor list
batch = [torch.rand(2, 3), torch.rand(2, 3)]
batch = trainer.accelerator.batch_to_device(batch, torch.device('cuda:0'))
assert batch[0].device.index == 0 and batch[0].type(
) == 'torch.cuda.FloatTensor'
assert batch[1].device.index == 0 and batch[1].type(
) == 'torch.cuda.FloatTensor'
# tensor list of lists
batch = [[torch.rand(2, 3), torch.rand(2, 3)]]
batch = trainer.accelerator.batch_to_device(batch, torch.device('cuda:0'))
assert batch[0][0].device.index == 0 and batch[0][0].type(
) == 'torch.cuda.FloatTensor'
assert batch[0][1].device.index == 0 and batch[0][1].type(
) == 'torch.cuda.FloatTensor'
# tensor dict
batch = [{'a': torch.rand(2, 3), 'b': torch.rand(2, 3)}]
batch = trainer.accelerator.batch_to_device(batch, torch.device('cuda:0'))
assert batch[0]['a'].device.index == 0 and batch[0]['a'].type(
) == 'torch.cuda.FloatTensor'
assert batch[0]['b'].device.index == 0 and batch[0]['b'].type(
) == 'torch.cuda.FloatTensor'
# tuple of tensor list and list of tensor dict
batch = ([torch.rand(2, 3) for _ in range(2)], [{
'a': torch.rand(2, 3),
'b': torch.rand(2, 3)
} for _ in range(2)])
batch = trainer.accelerator.batch_to_device(batch, torch.device('cuda:0'))
assert batch[0][0].device.index == 0 and batch[0][0].type(
) == 'torch.cuda.FloatTensor'
assert batch[1][0]['a'].device.index == 0
assert batch[1][0]['a'].type() == 'torch.cuda.FloatTensor'
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.accelerator.batch_to_device(batch, torch.device('cuda:0'))
assert batch[0].a.device.index == 0
assert batch[0].a.type() == 'torch.cuda.FloatTensor'
# non-Tensor that has `.to()` defined
class CustomBatchType:
def __init__(self):
self.a = torch.rand(2, 2)
def to(self, *args, **kwargs):
self.a = self.a.to(*args, **kwargs)
return self
batch = trainer.accelerator.batch_to_device(CustomBatchType(),
torch.device('cuda:0'))
assert batch.a.type() == 'torch.cuda.FloatTensor'
# torchtext.data.Batch
samples = [{
'text': 'PyTorch Lightning is awesome!',
'label': 0
}, {
'text': 'Please make it work with torchtext',
'label': 1
}]
text_field = Field()
label_field = LabelField()
fields = {'text': ('text', text_field), 'label': ('label', label_field)}
examples = [Example.fromdict(sample, fields) for sample in samples]
dataset = Dataset(examples=examples, fields=fields.values())
# Batch runs field.process() that numericalizes tokens, but it requires to build dictionary first
text_field.build_vocab(dataset)
label_field.build_vocab(dataset)
batch = Batch(data=examples, dataset=dataset)
batch = trainer.accelerator.batch_to_device(batch, torch.device('cuda:0'))
assert batch.text.type() == 'torch.cuda.LongTensor'
assert batch.label.type() == 'torch.cuda.LongTensor'
def test_single_gpu_batch_parse():
trainer = Trainer(accelerator="gpu", devices=1)
# non-transferrable types
primitive_objects = [None, {}, [], 1.0, "x", [None, 2], {"x": (1, 2), "y": None}]
for batch in primitive_objects:
data = trainer.strategy.batch_to_device(batch, torch.device("cuda:0"))
assert data == batch
# batch is just a tensor
batch = torch.rand(2, 3)
batch = trainer.strategy.batch_to_device(batch, torch.device("cuda:0"))
assert batch.device.index == 0 and batch.type() == "torch.cuda.FloatTensor"
# tensor list
batch = [torch.rand(2, 3), torch.rand(2, 3)]
batch = trainer.strategy.batch_to_device(batch, torch.device("cuda:0"))
assert batch[0].device.index == 0 and batch[0].type() == "torch.cuda.FloatTensor"
assert batch[1].device.index == 0 and batch[1].type() == "torch.cuda.FloatTensor"
# tensor list of lists
batch = [[torch.rand(2, 3), torch.rand(2, 3)]]
batch = trainer.strategy.batch_to_device(batch, torch.device("cuda:0"))
assert batch[0][0].device.index == 0 and batch[0][0].type() == "torch.cuda.FloatTensor"
assert batch[0][1].device.index == 0 and batch[0][1].type() == "torch.cuda.FloatTensor"
# tensor dict
batch = [{"a": torch.rand(2, 3), "b": torch.rand(2, 3)}]
batch = trainer.strategy.batch_to_device(batch, torch.device("cuda:0"))
assert batch[0]["a"].device.index == 0 and batch[0]["a"].type() == "torch.cuda.FloatTensor"
assert batch[0]["b"].device.index == 0 and batch[0]["b"].type() == "torch.cuda.FloatTensor"
# tuple of tensor list and list of tensor dict
batch = ([torch.rand(2, 3) for _ in range(2)], [{"a": torch.rand(2, 3), "b": torch.rand(2, 3)} for _ in range(2)])
batch = trainer.strategy.batch_to_device(batch, torch.device("cuda:0"))
assert batch[0][0].device.index == 0 and batch[0][0].type() == "torch.cuda.FloatTensor"
assert batch[1][0]["a"].device.index == 0
assert batch[1][0]["a"].type() == "torch.cuda.FloatTensor"
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.strategy.batch_to_device(batch, torch.device("cuda:0"))
assert batch[0].a.device.index == 0
assert batch[0].a.type() == "torch.cuda.FloatTensor"
# non-Tensor that has `.to()` defined
class CustomBatchType:
def __init__(self):
self.a = torch.rand(2, 2)
def to(self, *args, **kwargs):
self.a = self.a.to(*args, **kwargs)
return self
batch = trainer.strategy.batch_to_device(CustomBatchType(), torch.device("cuda:0"))
assert batch.a.type() == "torch.cuda.FloatTensor"
# torchtext.data.Batch
if not _TORCHTEXT_LEGACY:
return
samples = [
{"text": "PyTorch Lightning is awesome!", "label": 0},
{"text": "Please make it work with torchtext", "label": 1},
]
text_field = Field()
label_field = LabelField()
fields = {"text": ("text", text_field), "label": ("label", label_field)}
examples = [Example.fromdict(sample, fields) for sample in samples]
dataset = Dataset(examples=examples, fields=fields.values())
# Batch runs field.process() that numericalizes tokens, but it requires to build dictionary first
text_field.build_vocab(dataset)
label_field.build_vocab(dataset)
batch = Batch(data=examples, dataset=dataset)
with pytest.deprecated_call(match="The `torchtext.legacy.Batch` object is deprecated"):
batch = trainer.strategy.batch_to_device(batch, torch.device("cuda:0"))
assert batch.text.type() == "torch.cuda.LongTensor"
assert batch.label.type() == "torch.cuda.LongTensor"