传统方式需要10s，dat方式需要0.6s
import os
import time
import torch
import random
from common.coco_dataset import COCODataset
def gen_data(batch_size,data_path,target_path):
os.makedirs(target_path,exist_ok=True)
dataloader = torch.utils.data.DataLoader(COCODataset(data_path,
(352, 352),
is_training=False, is_scene=True),
batch_size=batch_size,
shuffle=False, num_workers=0, pin_memory=False,
drop_last=True) # DataLoader
start = time.time()
for step, samples in enumerate(dataloader):
images, labels, image_paths = samples["image"], samples["label"], samples["img_path"]
print("time", images.size(0), time.time() - start)
start = time.time()
# torch.save(samples,target_path+ '/' + str(step) + '.dat')
print(step)
def cat_100(target_path,batch_size=100):
paths = os.listdir(target_path)
li = [i for i in range(len(paths))]
random.shuffle(li)
images = []
labels = []
image_paths = []
start = time.time()
for i in range(len(paths)):
samples = torch.load(target_path + str(li[i]) + ".dat")
image, label, image_path = samples["image"], samples["label"], samples["img_path"]
images.append(image.cuda())
labels.append(label.cuda())
image_paths.append(image_path)
if i % batch_size == batch_size - 1:
images = torch.cat((images), 0)
print("time", images.size(0), time.time() - start)
images = []
labels = []
image_paths = []
start = time.time()
i += 1
if __name__ == '__main__':
os.environ["CUDA_VISIBLE_DEVICES"] = '3'
batch_size=320
# target_path='d:/test_1000/'
target_path='d:\img_2/'
data_path = r'D:\dataset\origin_all_datas\_2train'
gen_data(batch_size,data_path,target_path)
# get_data(target_path,batch_size)
# cat_100(target_path,batch_size)
这个读取数据也比较快：320 batch_size 450ms
def cat_100(target_path,batch_size=100):
paths = os.listdir(target_path)
li = [i for i in range(len(paths))]
random.shuffle(li)
images = []
labels = []
image_paths = []
start = time.time()
for i in range(len(paths)):
samples = torch.load(target_path + str(li[i]) + ".dat")
image, label, image_path = samples["image"], samples["label"], samples["img_path"]
images.append(image)#.cuda())
labels.append(label)#.cuda())
image_paths.append(image_path)
if i % batch_size < batch_size - 1:
i += 1
continue
i += 1
images = torch.cat(([image.cuda() for image in images]), 0)
print("time", images.size(0), time.time() - start)
images = []
labels = []
image_paths = []
start = time.time()