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"