Commit df7a70ab authored by novelailab's avatar novelailab

return identity

parent a0b66a4b
......@@ -82,21 +82,24 @@ class ShardedImageDataset(data.Dataset):
keys = [*range(key, key+self.bsz)]
tensors = self.executor.map(self.read_from_metadata_key, keys)
tensors = list(tensors)
tensors = torch.stack(tensors)
#make sure these operations are fast!
ids = [t[1] for t in tensors]
tensors = torch.stack([t[0] for t in tensors])
#####################################
if self.outer_transform:
tensors = self.outer_transform(tensors)
return tensors
def read_from_metadata_key(self, key):
offset, size, d_id = self.metadata[key]
offset, size, id = self.metadata[key]
data = self.mmap[offset:offset+size]
data = decode_jpeg(data)
data = torch.from_numpy(data).permute(2, 0, 1)
if self.inner_transform:
data = self.inner_transform(data)
return data
return data, id
class ImageDatasetBuilder():
def __init__(self, folder_path, name, threads=None):
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment