Ускорение Dataset.__getitems__ ⇐ Python

[Anonymous]

У меня есть модель с функцией пересылки, которая получает дополнительные параметры, например:

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

class MyModel(nn.Module):
...

def forward(self, interactions: torch.Tensor, user_features: Optional[torch.Tensor] = None):
"""
Where _N_ is the number of items

interactions (Tensor): Nx2
user_features (Tensor): Nx(number of features)
"""
...

По этой причине мой набор данных также возвращает словарь на основе DataFrame

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

class StackOverflowDataset(torch.utils.data.Dataset):
def __init__(self, data, user_features=None):
self._data = data
self._user_features = user_features

def __getitem__(self, idx):
if self._user_features is None:
return {'interactions': self._data[idx]}
else:
return {
'interactions': self._data[idx],
'user_features': self._user_features[self._data[idx]['user']]
}

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

Большая часть времени обучения тратится на __getitem__, что заставляет меня задуматься: не является ли использование необязательных аргументов в MyModel.forward плохой практикой? Я предполагаю, что большая часть времени тратится на игру с numpy поэлементно, а затем на преобразование его в тензоры PyTorch и перемещение на графический процессор. Есть ли способ заранее «предварительно обработать» все это, но при этом использовать Dataloader, который возвращает dict?
Похоже, что Datapipe также обрабатывает строку за строкой. Можно ли напрямую вернуть словарь со всем тензором «взаимодействий», а затем Dataloader разрежет его на части?



Подробнее здесь: https://stackoverflow.com/questions/790 ... t-getitems