mirror of
https://github.com/lucidrains/DALLE2-pytorch.git
synced 2025-12-19 17:54:20 +01:00
58892135d9
* Converted decoder trainer to use accelerate * Fixed issue where metric evaluation would hang on distributed mode * Implemented functional saving Loading still fails due to some issue with the optimizer * Fixed issue with loading decoders * Fixed issue with tracker config * Fixed issue with amp Updated logging to be more logical * Saving checkpoint now saves position in training as well Fixed an issue with running out of gpu space due to loading weights into the gpu twice * Fixed ema for distributed training * Fixed isue where get_pkg_version was reintroduced * Changed decoder trainer to upload config as a file Fixed issue where loading best would error
231 lines
12 KiB
Python
231 lines
12 KiB
Python
import os
|
|
import webdataset as wds
|
|
import torch
|
|
import numpy as np
|
|
import fsspec
|
|
import shutil
|
|
|
|
def get_shard(filename):
|
|
"""
|
|
Filenames with shards in them have a consistent structure that we can take advantage of
|
|
Standard structure: path/to/file/prefix_string_00001.ext
|
|
"""
|
|
try:
|
|
return filename.split("_")[-1].split(".")[0]
|
|
except ValueError:
|
|
raise RuntimeError(f"Could not find shard for filename {filename}")
|
|
|
|
def get_example_file(fs, path, file_format):
|
|
"""
|
|
Given a file system and a file extension, return the example file
|
|
"""
|
|
return fs.glob(os.path.join(path, f"*.{file_format}"))[0]
|
|
|
|
def embedding_inserter(samples, embeddings_url, index_width, handler=wds.handlers.reraise_exception):
|
|
"""Given a datum of {"__key__": str, "__url__": str, ...} adds the cooresponding embedding and yields"""
|
|
previous_tar_url = None
|
|
current_embeddings = None
|
|
# Get a reference to an abstract file system where the embeddings are stored
|
|
embeddings_fs, embeddings_path = fsspec.core.url_to_fs(embeddings_url)
|
|
example_embedding_file = get_example_file(embeddings_fs, embeddings_path, "npy")
|
|
example_embedding_shard = get_shard(example_embedding_file)
|
|
emb_shard_width = len(example_embedding_shard)
|
|
# Easier to get the basename without the shard once than search through for the correct file every time
|
|
embedding_file_basename = '_'.join(example_embedding_file.split("_")[:-1]) + "_"
|
|
|
|
def load_corresponding_embeds(tar_url):
|
|
"""Finds and reads the npy files that contains embeddings for the given webdataset tar"""
|
|
shard = int(tar_url.split("/")[-1].split(".")[0])
|
|
embedding_url = embedding_file_basename + str(shard).zfill(emb_shard_width) + '.npy'
|
|
with embeddings_fs.open(embedding_url) as f:
|
|
data = np.load(f)
|
|
return torch.from_numpy(data)
|
|
|
|
for sample in samples:
|
|
try:
|
|
tar_url = sample["__url__"]
|
|
key = sample["__key__"]
|
|
if tar_url != previous_tar_url:
|
|
# If the tar changed, we need to download new embeddings
|
|
# This means if we shuffle before inserting it will load many more files than we expect and be very inefficient.
|
|
previous_tar_url = tar_url
|
|
current_embeddings = load_corresponding_embeds(tar_url)
|
|
|
|
embedding_index = int(key[-index_width:])
|
|
embedding = current_embeddings[embedding_index]
|
|
# We need to check if this sample is nonzero. If it is, this embedding is not valid and we should continue to the next loop
|
|
if torch.count_nonzero(embedding) == 0:
|
|
raise RuntimeError(f"Webdataset had a sample, but no embedding was found. ImgShard: {key[:-index_width]} - Index: {key[-index_width:]}")
|
|
sample["npy"] = embedding
|
|
yield sample
|
|
except Exception as exn: # From wds implementation
|
|
if handler(exn):
|
|
continue
|
|
else:
|
|
break
|
|
insert_embedding = wds.filters.pipelinefilter(embedding_inserter)
|
|
|
|
def unassociated_shard_skipper(tarfiles, embeddings_url, handler=wds.handlers.reraise_exception):
|
|
"""Finds if the is a corresponding embedding for the tarfile at { url: [URL] }"""
|
|
embeddings_fs, embeddings_path = fsspec.core.url_to_fs(embeddings_url)
|
|
embedding_files = embeddings_fs.ls(embeddings_path)
|
|
get_embedding_shard = lambda embedding_file: int(embedding_file.split("_")[-1].split(".")[0])
|
|
embedding_shards = set([get_embedding_shard(filename) for filename in embedding_files]) # Sets have O(1) check for member
|
|
|
|
get_tar_shard = lambda tar_file: int(tar_file.split("/")[-1].split(".")[0])
|
|
for tarfile in tarfiles:
|
|
try:
|
|
webdataset_shard = get_tar_shard(tarfile["url"])
|
|
# If this shard has an associated embeddings file, we pass it through. Otherwise we iterate until we do have one
|
|
if webdataset_shard in embedding_shards:
|
|
yield tarfile
|
|
except Exception as exn: # From wds implementation
|
|
if handler(exn):
|
|
continue
|
|
else:
|
|
break
|
|
|
|
skip_unassociated_shards = wds.filters.pipelinefilter(unassociated_shard_skipper)
|
|
|
|
def verify_keys(samples, handler=wds.handlers.reraise_exception):
|
|
"""
|
|
Requires that both the image and embedding are present in the sample
|
|
This is important to do as a user may forget they do not have embeddings in their webdataset and neglect to add them using the embedding_folder_url parameter.
|
|
"""
|
|
for sample in samples:
|
|
try:
|
|
assert "jpg" in sample, f"Sample {sample['__key__']} missing image"
|
|
assert "npy" in sample, f"Sample {sample['__key__']} missing embedding. Did you set embedding_folder_url?"
|
|
yield sample
|
|
except Exception as exn: # From wds implementation
|
|
if handler(exn):
|
|
continue
|
|
else:
|
|
break
|
|
|
|
class ImageEmbeddingDataset(wds.DataPipeline, wds.compat.FluidInterface):
|
|
"""
|
|
A fluid interface wrapper for DataPipline that returns image embedding pairs
|
|
Reads embeddings as npy files from the webdataset if they exist. If embedding_folder_url is set, they will be inserted in from the alternate source.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
urls,
|
|
embedding_folder_url=None,
|
|
index_width=None,
|
|
img_preproc=None,
|
|
extra_keys=[],
|
|
handler=wds.handlers.reraise_exception,
|
|
resample=False,
|
|
shuffle_shards=True
|
|
):
|
|
"""
|
|
Modeled directly off of the WebDataset constructor
|
|
|
|
:param urls: A url pointing to the tar files of the webdataset formatted as /path/to/webdataset/{0000..9999}.tar
|
|
:param embedding_folder_url: Required if webdataset does not contain embeddings. A url pointing to the npy files of the embeddings. Should have the same number of shards as the webdataset.
|
|
Webdataset image keys should align with the index of the embedding. This means missing image indices must have a corresponding embedding of all zeros.
|
|
:param index_width: The number of digits in the index. This is used to align the embedding index with the image index.
|
|
For example, if a file in the webdataset shard 3 is named 0003039.jpg, we know the shard is 4 digits and the last 3 digits are the index_width.
|
|
:param img_preproc: This function is run on the img before it is batched and returned. Useful for data augmentation or converting to torch tensor.
|
|
:param handler: A webdataset handler.
|
|
:param resample: If true, resample webdataset shards with replacement. You need to set your own epoch size if this is true since it will resample infinitely.
|
|
:param shuffle_shards: If true, shuffle the shards before resampling. This cannot be true if resample is true.
|
|
|
|
|
|
"""
|
|
super().__init__()
|
|
keys = ["jpg", "npy"] + extra_keys
|
|
self.key_map = {key: i for i, key in enumerate(keys)}
|
|
self.resampling = resample
|
|
self.img_preproc = img_preproc
|
|
# If s3, check if s3fs is installed and s3cmd is installed and check if the data is piped instead of straight up
|
|
if (isinstance(urls, str) and "s3:" in urls) or (isinstance(urls, list) and any(["s3:" in url for url in urls])):
|
|
# Then this has an s3 link for the webdataset and we need extra packages
|
|
if shutil.which("s3cmd") is None:
|
|
raise RuntimeError("s3cmd is required for s3 webdataset")
|
|
if "s3:" in embedding_folder_url:
|
|
# Then the embeddings are being loaded from s3 and fsspec requires s3fs
|
|
try:
|
|
import s3fs
|
|
except ImportError:
|
|
raise RuntimeError("s3fs is required to load embeddings from s3")
|
|
# Add the shardList and randomize or resample if requested
|
|
if resample:
|
|
assert not shuffle_shards, "Cannot both resample and shuffle"
|
|
self.append(wds.ResampledShards(urls))
|
|
else:
|
|
self.append(wds.SimpleShardList(urls))
|
|
if shuffle_shards:
|
|
self.append(wds.filters.shuffle(1000))
|
|
|
|
if embedding_folder_url is not None:
|
|
# There may be webdataset shards that do not have a embedding shard associated with it. If we do not skip these, they would cause issues.
|
|
self.append(skip_unassociated_shards(embeddings_url=embedding_folder_url, handler=handler))
|
|
|
|
self.append(wds.tarfile_to_samples(handler=handler))
|
|
self.append(wds.decode("pilrgb", handler=handler))
|
|
if embedding_folder_url is not None:
|
|
# Then we are loading embeddings for a remote source
|
|
assert index_width is not None, "Reading embeddings separately requires index width length to be given"
|
|
self.append(insert_embedding(embeddings_url=embedding_folder_url, index_width=index_width, handler=handler))
|
|
self.append(verify_keys)
|
|
# Apply preprocessing
|
|
self.append(wds.map(self.preproc))
|
|
self.append(wds.to_tuple(*keys))
|
|
|
|
def preproc(self, sample):
|
|
"""Applies the preprocessing for images"""
|
|
if self.img_preproc is not None:
|
|
sample["jpg"] = self.img_preproc(sample["jpg"])
|
|
return sample
|
|
|
|
def create_image_embedding_dataloader(
|
|
tar_url,
|
|
num_workers,
|
|
batch_size,
|
|
embeddings_url=None,
|
|
index_width=None,
|
|
shuffle_num = None,
|
|
shuffle_shards = True,
|
|
resample_shards = False,
|
|
img_preproc=None,
|
|
extra_keys=[],
|
|
handler=wds.handlers.reraise_exception#warn_and_continue
|
|
):
|
|
"""
|
|
Convenience function to create an image embedding dataseta and dataloader in one line
|
|
|
|
:param tar_url: A url pointing to the tar files of the webdataset formatted as /path/to/webdataset/{0000..9999}.tar
|
|
:param num_workers: The number of workers to use for the dataloader
|
|
:param batch_size: The batch size to use for the dataloader
|
|
:param embeddings_url: Required if webdataset does not contain embeddings. A url pointing to the npy files of the embeddings. Should have the same number of shards as the webdataset.
|
|
Webdataset image keys should align with the index of the embedding. This means missing image indices must have a corresponding embedding of all zeros.
|
|
:param index_width: The number of digits in the index. This is used to align the embedding index with the image index.
|
|
For example, if a file in the webdataset shard 3 is named 0003039.jpg, we know the shard is 4 digits and the last 3 digits are the index_width.
|
|
:param shuffle_num: If not None, shuffle the dataset with this size buffer after sampling.
|
|
:param shuffle_shards: If true, shuffle the shards before sampling. This cannot be true if resample is true.
|
|
:param resample_shards: If true, resample webdataset shards with replacement. You need to set your own epoch size if this is true since it will resample infinitely.
|
|
:param handler: A webdataset handler.
|
|
"""
|
|
ds = ImageEmbeddingDataset(
|
|
tar_url,
|
|
embeddings_url,
|
|
index_width=index_width,
|
|
shuffle_shards=shuffle_shards,
|
|
resample=resample_shards,
|
|
extra_keys=extra_keys,
|
|
img_preproc=img_preproc,
|
|
handler=handler
|
|
)
|
|
if shuffle_num is not None and shuffle_num > 0:
|
|
ds.shuffle(1000)
|
|
return wds.WebLoader(
|
|
ds,
|
|
num_workers=num_workers,
|
|
batch_size=batch_size,
|
|
prefetch_factor=2, # This might be good to have high so the next npy file is prefetched
|
|
pin_memory=True,
|
|
shuffle=False
|
|
) |