Updated number of batches calculation.
@@ -89,6 +89,29 @@ class BERTSequenceClassifier:
)
self.model = move_to_device(self.model, device, num_gpus)
+ token_ids_tensor = torch.tensor(token_ids, dtype=torch.long)
+ input_mask_tensor = torch.tensor(input_mask, dtype=torch.long)
+ labels_tensor = torch.tensor(labels, dtype=torch.long)
+
+ if token_type_ids:
+ token_type_ids_tensor = torch.tensor(
+ token_type_ids, dtype=torch.long
+ )
+ train_dataset = TensorDataset(
+ token_ids_tensor,
+ input_mask_tensor,
+ token_type_ids_tensor,
+ labels_tensor,
+ )
+ else:
+ train_dataset = TensorDataset(
+ token_ids_tensor, input_mask_tensor, labels_tensor
+ )
+ train_sampler = RandomSampler(train_dataset)
+
+ train_dataloader = DataLoader(
+ train_dataset, sampler=train_sampler, batch_size=batch_size
+ )
# define optimizer and model parameters
param_optimizer = list(self.model.named_parameters())
no_decay = ["bias", "LayerNorm.bias", "LayerNorm.weight"]
@@ -111,8 +134,7 @@ class BERTSequenceClassifier:
},
]
- num_examples = len(token_ids)
- num_batches = int(num_examples / batch_size)
+ num_batches = len(train_dataloader)
num_train_optimization_steps = num_batches * num_epochs
if warmup_proportion is None:
@@ -131,30 +153,6 @@ class BERTSequenceClassifier:
# train
self.model.train() # training mode
- token_ids_tensor = torch.tensor(token_ids, dtype=torch.long)
- input_mask_tensor = torch.tensor(input_mask, dtype=torch.long)
- labels_tensor = torch.tensor(labels, dtype=torch.long)
-
- if token_type_ids:
- token_type_ids_tensor = torch.tensor(
- token_type_ids, dtype=torch.long
- )
- train_dataset = TensorDataset(
- token_ids_tensor,
- input_mask_tensor,
- token_type_ids_tensor,
- labels_tensor,
- )
- else:
- train_dataset = TensorDataset(
- token_ids_tensor, input_mask_tensor, labels_tensor
- )
- train_sampler = RandomSampler(train_dataset)
-
- train_dataloader = DataLoader(
- train_dataset, sampler=train_sampler, batch_size=batch_size
- )
-
for epoch in range(num_epochs):
training_loss = 0
for i, batch in enumerate(
