|
|
|
|
|
import os |
|
|
import json |
|
|
import tempfile |
|
|
import numpy as np |
|
|
import torch |
|
|
import time |
|
|
import subprocess |
|
|
import torch.distributed as dist |
|
|
|
|
|
|
|
|
def allreduce(x, average): |
|
|
if mpi_size() > 1: |
|
|
dist.all_reduce(x, dist.ReduceOp.SUM) |
|
|
return x / mpi_size() if average else x |
|
|
|
|
|
|
|
|
def get_cpu_stats_over_ranks(stat_dict): |
|
|
keys = sorted(stat_dict.keys()) |
|
|
allreduced = allreduce(torch.stack([torch.as_tensor(stat_dict[k]).detach().cuda().float() for k in keys]), average=True).cpu() |
|
|
return {k: allreduced[i].item() for (i, k) in enumerate(keys)} |
|
|
|
|
|
|
|
|
class Hyperparams(dict): |
|
|
def __getattr__(self, attr): |
|
|
try: |
|
|
return self[attr] |
|
|
except KeyError: |
|
|
return None |
|
|
|
|
|
def __setattr__(self, attr, value): |
|
|
self[attr] = value |
|
|
|
|
|
|
|
|
def logger(log_prefix): |
|
|
'Prints the arguments out to stdout, .txt, and .jsonl files' |
|
|
|
|
|
jsonl_path = f'{log_prefix}.jsonl' |
|
|
txt_path = f'{log_prefix}.txt' |
|
|
|
|
|
def log(*args, pprint=False, **kwargs): |
|
|
if mpi_rank() != 0: |
|
|
return |
|
|
t = time.ctime() |
|
|
argdict = {'time': t} |
|
|
if len(args) > 0: |
|
|
argdict['message'] = ' '.join([str(x) for x in args]) |
|
|
argdict.update(kwargs) |
|
|
|
|
|
txt_str = [] |
|
|
args_iter = sorted(argdict) if pprint else argdict |
|
|
for k in args_iter: |
|
|
val = argdict[k] |
|
|
if isinstance(val, np.ndarray): |
|
|
val = val.tolist() |
|
|
elif isinstance(val, np.integer): |
|
|
val = int(val) |
|
|
elif isinstance(val, np.floating): |
|
|
val = float(val) |
|
|
argdict[k] = val |
|
|
if isinstance(val, float): |
|
|
val = f'{val:.5f}' |
|
|
txt_str.append(f'{k}: {val}') |
|
|
txt_str = ', '.join(txt_str) |
|
|
|
|
|
if pprint: |
|
|
json_str = json.dumps(argdict, sort_keys=True) |
|
|
txt_str = json.dumps(argdict, sort_keys=True, indent=4) |
|
|
else: |
|
|
json_str = json.dumps(argdict) |
|
|
|
|
|
print(txt_str, flush=True) |
|
|
|
|
|
with open(txt_path, "a+") as f: |
|
|
print(txt_str, file=f, flush=True) |
|
|
with open(jsonl_path, "a+") as f: |
|
|
print(json_str, file=f, flush=True) |
|
|
|
|
|
return log |
|
|
|
|
|
|
|
|
def maybe_download(path, filename=None): |
|
|
'''If a path is a gsutil path, download it and return the local link, |
|
|
otherwise return link''' |
|
|
if not path.startswith('gs://'): |
|
|
return path |
|
|
if filename: |
|
|
local_dest = f'/tmp/' |
|
|
out_path = f'/tmp/{filename}' |
|
|
if os.path.isfile(out_path): |
|
|
return out_path |
|
|
subprocess.check_output(['gsutil', '-m', 'cp', '-R', path, out_path]) |
|
|
return out_path |
|
|
else: |
|
|
local_dest = tempfile.mkstemp()[1] |
|
|
subprocess.check_output(['gsutil', '-m', 'cp', path, local_dest]) |
|
|
return local_dest |
|
|
|
|
|
|
|
|
def tile_images(images, d1=4, d2=4, border=1): |
|
|
id1, id2, c = images[0].shape |
|
|
out = np.ones([d1 * id1 + border * (d1 + 1), |
|
|
d2 * id2 + border * (d2 + 1), |
|
|
c], dtype=np.uint8) |
|
|
out *= 255 |
|
|
if len(images) != d1 * d2: |
|
|
raise ValueError('Wrong num of images') |
|
|
for imgnum, im in enumerate(images): |
|
|
num_d1 = imgnum // d2 |
|
|
num_d2 = imgnum % d2 |
|
|
start_d1 = num_d1 * id1 + border * (num_d1 + 1) |
|
|
start_d2 = num_d2 * id2 + border * (num_d2 + 1) |
|
|
out[start_d1:start_d1 + id1, start_d2:start_d2 + id2, :] = im |
|
|
return out |
|
|
|
|
|
|
|
|
def mpi_size(): |
|
|
return MPI.COMM_WORLD.Get_size() |
|
|
|
|
|
|
|
|
def mpi_rank(): |
|
|
return MPI.COMM_WORLD.Get_rank() |
|
|
|
|
|
|
|
|
def num_nodes(): |
|
|
nn = mpi_size() |
|
|
if nn % 8 == 0: |
|
|
return nn // 8 |
|
|
return nn // 8 + 1 |
|
|
|
|
|
|
|
|
def gpus_per_node(): |
|
|
size = mpi_size() |
|
|
if size > 1: |
|
|
return max(size // num_nodes(), 1) |
|
|
return 1 |
|
|
|
|
|
|
|
|
def local_mpi_rank(): |
|
|
return mpi_rank() % gpus_per_node() |
|
|
|
