Question

我正在使用torch.utils.dataloaders创建具有几千个图像的数据加载器（阶段0），并对32个图像批次进行正向传递（阶段1）。到目前为止，我已经能够在推理阶段（阶段2）优化GPU的使用和内存的使用。在此期间，我将这些批次的输出收集到python列表中，然后在单独的后处理步骤中，将所有这些输出打包到一个文件中（压缩它们之前）。

我注意到，每30k图像，第0阶段花费的时间可忽略不计，而对于第1阶段，向前通过则花费120秒（每秒约250张图像），而包装花费的时间则成比例（约100秒）。在打包期间，GPU保持空闲状态，只有在将输出保存到磁盘之后，我们才准备下一个数据加载器。如何避免这种情况？有没有更好的方法来改善后处理步骤？

注意：数百万张图像需要重复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)

优化推断后的输出包装时间

0 个答案: