Added changes to allow test loader creation with Sequential
@@ -136,11 +136,12 @@ class BERTSequenceClassifier:
return optimizer
- def create_data_loader(self, dataset, batch_size=32, **kwargs):
+ def create_data_loader(self, dataset, batch_size=32, mode="train", **kwargs):
"""
Method to create a data loader for a given Tensor dataset.
Args:
+ mode(str): Mode for creating data loader. Could be train or test.
dataset(torch.utils.data.Dataset): A Tensor dataset.
batch_size(int): Batch size.
@@ -149,15 +150,17 @@ class BERTSequenceClassifier:
"""
- if self.use_distributed:
- train_sampler = torch.utils.data.distributed.DistributedSampler(
+ if mode == "test":
+ sampler = torch.utils.data.sampler.SequentialSampler(dataset)
+ elif self.use_distributed:
+ sampler = torch.utils.data.distributed.DistributedSampler(
dataset, num_replicas=hvd.size(), rank=hvd.rank()
)
else:
- train_sampler = torch.utils.data.RandomSampler(dataset)
+ sampler = torch.utils.data.RandomSampler(dataset)
data_loader = torch.utils.data.DataLoader(
- dataset, batch_size=batch_size, sampler=train_sampler, **kwargs
+ dataset, batch_size=batch_size, sampler=sampler, **kwargs
)
return data_loader
@@ -287,6 +290,9 @@ class BERTSequenceClassifier:
def predict(self, test_loader, num_gpus=None, probabilities=False):
"""
+ Method to predict the results on the test loader. Only evaluates for non distributed
+ workload on the head node in a distributed setup.
+
Args:
test_loader(torch Dataloader): Torch Dataloader created from Torch Dataset
num_gpus (int, optional): The number of gpus to use.
@@ -300,8 +306,6 @@ class BERTSequenceClassifier:
1darray, dict(1darray, 1darray, ndarray): Predicted classes and target labels or
a dictionary with classes, target labels, probabilities) if probabilities is True.
"""
-
- if hvd.rank() == 0:
device = get_device("cpu" if num_gpus == 0 else "gpu")
self.model = move_to_device(self.model, device, num_gpus)
@@ -320,7 +324,10 @@ class BERTSequenceClassifier:
y_batch = data["labels"]
token_type_ids_batch = None
- if "token_type_ids" in data and data["token_type_ids"] is not None:
+ if (
+ "token_type_ids" in data
+ and data["token_type_ids"] is not None
+ ):
token_type_ids_batch = data["token_type_ids"]
token_type_ids_batch = token_type_ids_batch.cuda()
