Update sequence_classification.py
import numpy as np
import torch
-import torch.nn as nn
from torch.utils.data import DataLoader, RandomSampler, SequentialSampler
from torch.utils.data.distributed import DistributedSampler
-
from transformers.modeling_bert import (
BERT_PRETRAINED_MODEL_ARCHIVE_MAP,
BertForSequenceClassification,
@@ -23,11 +21,9 @@ from transformers.modeling_xlnet import (
XLNET_PRETRAINED_MODEL_ARCHIVE_MAP,
XLNetForSequenceClassification,
)
-from utils_nlp.common.pytorch_utils import get_device
from utils_nlp.models.transformers.common import MAX_SEQ_LEN, TOKENIZER_CLASS, Transformer
from utils_nlp.models.transformers.datasets import SCDataSet, SPCDataSet
-
MODEL_CLASS = {}
MODEL_CLASS.update({k: BertForSequenceClassification for k in BERT_PRETRAINED_MODEL_ARCHIVE_MAP})
MODEL_CLASS.update(
@@ -285,16 +281,9 @@ class SequenceClassifier(Transformer):
seed (int, optional): Random seed used to improve reproducibility. Defaults to None.
"""
- device, num_gpus = get_device(num_gpus=num_gpus, local_rank=local_rank)
- if isinstance(self.model, nn.DataParallel):
- self.model.module.to(device)
- else:
- self.model.to(device)
-
super().fine_tune(
train_dataloader=train_dataloader,
get_inputs=Processor.get_inputs,
- device=device,
n_gpu=num_gpus,
num_train_epochs=num_epochs,
weight_decay=weight_decay,
@@ -319,15 +308,13 @@ class SequenceClassifier(Transformer):
Returns
1darray: numpy array of predicted label indices.
"""
- device, num_gpus = get_device(num_gpus=num_gpus, local_rank=-1)
preds = list(
super().predict(
eval_dataloader=eval_dataloader,
get_inputs=Processor.get_inputs,
- device=device,
- verbose=verbose,
n_gpu=num_gpus,
+ verbose=verbose,
)
)
preds = np.concatenate(preds)
