Почему мой PyTorch DataLoader использует только одно ядро ​​ЦП, несмотря на настройку num_workers>1? ⇐ Python

[Anonymous]

Я пытаюсь настроить BERT для задачи классификации по нескольким меткам (ядовитые комментарии Jigsaw). Я создал собственный набор данных и DataLoader следующим образом:

Код: Выделить всё

class CustomDataSet(Dataset):

def __init__(self,
features: np.ndarray,
labels: np.ndarray,
token_max: int,
tokenizer):

self.features = features
self.labels = labels
self.tokenizer = tokenizer
self.token_max = token_max

def __len__(self):
return len(self.features)

def __getitem__(self,
index: int):
comment_id, comment_text = self.features[index]
labels = self.labels[index]

encoding = self.tokenizer.encode_plus(
comment_text,
add_special_tokens=True,
max_length=self.token_max,
return_token_type_ids=False,
padding='max_length',
truncation=True,
return_attention_mask=True,
return_tensors='pt')

return dict(
comment_text=comment_text,
comment_id=comment_id,
input_ids=encoding['input_ids'].squeeze(0),
attention_mask=encoding['attention_mask'].squeeze(0),
labels=torch.Tensor(labels))

и я использую следующий токенизатор:

Код: Выделить всё

tokenizer = BertTokenizer.from_pretrained('bert-base-cased')

затем я создаю свой набор данных, используя пользовательский класс и соответствующий загрузчик данных:

Код: Выделить всё

train_dataset = CustomDataSet(X_train, y_train, tokenizer=tokenizer, token_max=256)
train_loader = DataLoader(
train_dataset, batch_size=32, shuffle=True, pin_memory=True, num_workers=16, persistent_workers=False
)

Моя модель следующая:

Код: Выделить всё

class MultiLabelBERT(torch.nn.Module):
def __init__(self, num_labels):
super(MultiLabelBERT, self).__init__()
self.bert = BertModel.from_pretrained("bert-base-uncased", torch_dtype=torch.float16, attn_implementation="sdpa")
self.classifier = torch.nn.Linear(self.bert.config.hidden_size, num_labels)
self.classifier = self.classifier.to(torch.float16)

def forward(self, input_ids, attention_mask):
outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
pooled_output = outputs.pooler_output
logits = self.classifier(pooled_output)
return logits

first_BERT = MultiLabelBERT(6)

Сначала я хочу перебрать все пакеты моего обучающего набора и выполнить прямой проход:

Код: Выделить всё

for batch_idx, item  in enumerate(train_loader):
input_ids = item['input_ids'].to(self.device)
attention_mask = item['attention_mask'].to(self.device)
labels = item['labels'].to(self.device)
logits = first_BERT(input_ids=input_ids,
attention_mask=attention_mask)

Хотя я установил num_workers=16 в своем DataLoader, он использует только одно ядро ​​ЦП для загрузки данных в мой графический процессор. Это существенно замедляет процесс. Вот что я пробовал:

• Уменьшить размер пакета.
• Уменьшить количество токенов (token_max).
• Предварительная токенизация всего набора данных, чтобы гарантировать, что токенизатор не вызывает проблем с узким местом.
  Когда я комментирую прямой проход, DataLoader использует все рабочие процессы ЦП, как и ожидалось. Однако при прямом проходе процесс кажется узким местом.
  У меня следующая настройка графического процессора:

Код: Выделить всё

+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 550.120                Driver Version: 550.120        CUDA Version: 12.4     |
|-----------------------------------------+------------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
|                                         |                        |               MIG M. |
|=========================================+========================+======================|
|   0  NVIDIA GeForce RTX 4060 ...     Off |   00000000:01:00.0 Off |                  N/A |
| N/A   40C    P0            588W /  115W |       9MiB /   8188MiB |      0%      Default |
|                                         |                        |                  N/A |
+-----------------------------------------+------------------------+----------------------+

+-----------------------------------------------------------------------------------------+
| Processes:                                                                              |
|  GPU   GI   CI        PID   Type   Process name                              GPU Memory |
|        ID   ID                                                               Usage      |
|=========================================================================================|
|    0   N/A  N/A      1285      G   /usr/lib/xorg/Xorg                              4MiB |

Кто-нибудь знает, что может вызвать это?

Подробнее здесь: https://stackoverflow.com/questions/793 ... um-workers