Spaces:
Sleeping
Sleeping
File size: 3,084 Bytes
e75a247 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 |
import glob
import numpy as np
import os
import shutil
from src.logger.logger import _logger
def to_filelist(args, mode="train"):
if mode == "train":
flist = args.data_train
elif mode == "val":
flist = args.data_val
elif mode == "test":
flist = args.data_test
else:
raise NotImplementedError("Invalid mode %s" % mode)
print(flist)
# keyword-based: 'a:/path/to/a b:/path/to/b'
file_dict = {}
for f in flist:
if ":" in f:
name, fp = f.split(":")
else:
name, fp = "_", f
files = glob.glob(fp)
if name in file_dict:
file_dict[name] += files
else:
file_dict[name] = files
# sort files
for name, files in file_dict.items():
file_dict[name] = sorted(files)
if args.local_rank is not None:
if mode == "train":
local_world_size = int(os.environ["LOCAL_WORLD_SIZE"])
new_file_dict = {}
for name, files in file_dict.items():
new_files = files[args.local_rank :: local_world_size]
assert len(new_files) > 0
np.random.shuffle(new_files)
new_file_dict[name] = new_files
file_dict = new_file_dict
if args.copy_inputs:
import tempfile
tmpdir = tempfile.mkdtemp()
if os.path.exists(tmpdir):
shutil.rmtree(tmpdir)
new_file_dict = {name: [] for name in file_dict}
for name, files in file_dict.items():
for src in files:
dest = os.path.join(tmpdir, src.lstrip("/"))
if not os.path.exists(os.path.dirname(dest)):
os.makedirs(os.path.dirname(dest), exist_ok=True)
shutil.copy2(src, dest)
_logger.info("Copied file %s to %s" % (src, dest))
new_file_dict[name].append(dest)
if len(files) != len(new_file_dict[name]):
_logger.error(
"Only %d/%d files copied for %s file group %s",
len(new_file_dict[name]),
len(files),
mode,
name,
)
file_dict = new_file_dict
filelist = sum(file_dict.values(), [])
assert len(filelist) == len(set(filelist))
return file_dict, filelist
def clear_empty_paths(dir):
# clear the dirs in this folder that are empty (i.e. don't have any files or folders in them)
for f in os.listdir(dir):
if not os.path.isdir(os.path.join(dir, f)):
continue
if not os.listdir(os.path.join(dir, f)):
shutil.rmtree(os.path.join(dir, f))
_logger.info("Removed empty path %s" % f)
import io
import torch
import pickle
class CPU_Unpickler(pickle.Unpickler):
def find_class(self, module, name):
if module == 'torch.storage' and name == '_load_from_bytes':
return lambda b: torch.load(io.BytesIO(b), map_location='cpu')
else:
return super().find_class(module, name)
|