Я использую torch.utils.dataloaders, чтобы создать загрузчик данных (фаза 0) с несколькими тысячами изображений и выполнить прямую передачу (фаза 1) для пакетов из 32 изображений.До сих пор я был в состоянии оптимизировать использование графического процессора и памяти на этапе вывода (этап 2).В течение этого периода я собираю выходные данные этих пакетов в списки Python, а затем на отдельном этапе последующей обработки упаковываю все эти выходные данные в файл (перед их сжатием).
Я заметил, что на 30 тыс. ИзображенийФаза 0 занимает незначительное время, тогда как для Фазы 1 прямой проход занимает 120 секунд (около 250 изображений в секунду), а упаковка занимает пропорциональное количество времени (около 100 секунд).Во время упаковки графический процессор остается бездействующим, и только после сохранения выходных данных на диск мы готовим следующий загрузчик данных.Как я могу избежать этого?Есть ли лучший способ улучшить этап постобработки?
Примечание. Фазы 0-3 необходимо повторить для миллионов изображений.
Вот примерный скелет моего кода:
### PHASE 0 : PREPARING DATA
dataset = CustomDataset(items)
data_loader = torch.utils.data.DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=NUM_PREPROCESS_WORKERS, drop_last=False, collate_fn=filtered_collate_fn)
cids, features, meta = [], [], []
### PHASE 1 : FORWARD PASS #####
# note each of *_batch is a tuple
for batch_idx, (id_batch, img_batch, tag_batch) in enumerate(data_loader):
try:
feat_batch = self.model.forward_pass(img_batch)
cids = [*cids, *list(id_batch)]
features = [*features, *list(feat_batch)]
meta = [*meta, *list(meta_batch)]
### PHASE 2 : PACKAGING OUTPUTS of PHASE 2 #####
filename_pattern = 'chunk-{}.json'.format(get_random_string_with_timestamp())
items = []
for cid, feature, tag in zip(cids, features, meta):
json_obj = { "cid" : str(cid), "feature": feature, "tag" : str(tag) }
items.append(json.dumps(json_obj))
file_data = '\n'.join(items)
f.write(file_data)