|
|
import torch |
|
|
import collections |
|
|
import logging |
|
|
import omegaconf |
|
|
import wandb |
|
|
import datetime |
|
|
import glob |
|
|
import os |
|
|
import json |
|
|
|
|
|
from PIL import Image |
|
|
|
|
|
|
|
|
class BaseTimer: |
|
|
def __init__(self): |
|
|
self.start = torch.cuda.Event(enable_timing=True) |
|
|
self.end = torch.cuda.Event(enable_timing=True) |
|
|
self.start.record() |
|
|
|
|
|
def stop(self): |
|
|
self.end.record() |
|
|
torch.cuda.synchronize() |
|
|
return self.start.elapsed_time(self.end) / 1000 |
|
|
|
|
|
class Timer: |
|
|
def __init__(self, info=None, log_event=None): |
|
|
self.info = info |
|
|
self.log_event = log_event |
|
|
|
|
|
def __enter__(self): |
|
|
self.start = torch.cuda.Event(enable_timing=True) |
|
|
self.end = torch.cuda.Event(enable_timing=True) |
|
|
self.start.record() |
|
|
return self |
|
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb): |
|
|
self.end.record() |
|
|
torch.cuda.synchronize() |
|
|
self.duration = self.start.elapsed_time(self.end) / 1000 |
|
|
if self.info: |
|
|
self.info[f"duration/{self.log_event}"] = self.duration |
|
|
|
|
|
|
|
|
class _StreamingMean: |
|
|
def __init__(self, val=None, counts=None): |
|
|
if val is None: |
|
|
self.mean = 0.0 |
|
|
self.counts = 0 |
|
|
else: |
|
|
if isinstance(val, torch.Tensor): |
|
|
val = val.data.cpu().numpy() |
|
|
self.mean = val |
|
|
if counts is not None: |
|
|
self.counts = counts |
|
|
else: |
|
|
self.counts = 1 |
|
|
|
|
|
def update(self, mean, counts=1): |
|
|
if isinstance(mean, torch.Tensor): |
|
|
mean = mean.data.cpu().numpy() |
|
|
elif isinstance(mean, _StreamingMean): |
|
|
mean, counts = mean.mean, mean.counts * counts |
|
|
assert counts >= 0 |
|
|
if counts == 0: |
|
|
return |
|
|
total = self.counts + counts |
|
|
self.mean = self.counts / total * self.mean + counts / total * mean |
|
|
self.counts = total |
|
|
|
|
|
def __add__(self, other): |
|
|
new = self.__class__(self.mean, self.counts) |
|
|
if isinstance(other, _StreamingMean): |
|
|
if other.counts == 0: |
|
|
return new |
|
|
else: |
|
|
new.update(other.mean, other.counts) |
|
|
else: |
|
|
new.update(other) |
|
|
return new |
|
|
|
|
|
|
|
|
class StreamingMeans(collections.defaultdict): |
|
|
def __init__(self): |
|
|
super().__init__(_StreamingMean) |
|
|
|
|
|
def __setitem__(self, key, value): |
|
|
if isinstance(value, _StreamingMean): |
|
|
super().__setitem__(key, value) |
|
|
else: |
|
|
super().__setitem__(key, _StreamingMean(value)) |
|
|
|
|
|
def update(self, *args, **kwargs): |
|
|
for_update = dict(*args, **kwargs) |
|
|
for k, v in for_update.items(): |
|
|
self[k].update(v) |
|
|
|
|
|
def to_dict(self, prefix=""): |
|
|
return dict((prefix + k, v.mean) for k, v in self.items()) |
|
|
|
|
|
def to_str(self): |
|
|
return ", ".join([f"{k} = {v:.3f}" for k, v in self.to_dict().items()]) |
|
|
|
|
|
|
|
|
class ConsoleLogger: |
|
|
def __init__(self, name): |
|
|
self.logger = logging.getLogger(name) |
|
|
self.logger.handlers = [] |
|
|
self.logger.setLevel(logging.INFO) |
|
|
log_formatter = logging.Formatter( |
|
|
"%(asctime)s %(message)s", datefmt="%Y-%m-%d %H:%M:%S" |
|
|
) |
|
|
console_handler = logging.StreamHandler() |
|
|
console_handler.setFormatter(log_formatter) |
|
|
self.logger.addHandler(console_handler) |
|
|
|
|
|
self.logger.propagate = False |
|
|
|
|
|
@staticmethod |
|
|
def format_info(info): |
|
|
if not info: |
|
|
return str(info) |
|
|
log_groups = collections.defaultdict(dict) |
|
|
for k, v in info.to_dict().items(): |
|
|
prefix, suffix = k.split("/", 1) |
|
|
log_groups[prefix][suffix] = f"{v:.3f}" if isinstance(v, float) else str(v) |
|
|
formatted_info = "" |
|
|
max_group_size = len(max(log_groups, key=len)) + 2 |
|
|
max_k_size = max([len(max(g, key=len)) for g in log_groups.values()]) + 1 |
|
|
max_v_size = ( |
|
|
max([len(max(g.values(), key=len)) for g in log_groups.values()]) + 1 |
|
|
) |
|
|
for group, group_info in log_groups.items(): |
|
|
group_str = [ |
|
|
f"{k:<{max_k_size}}={v:>{max_v_size}}" for k, v in group_info.items() |
|
|
] |
|
|
max_g_size = len(max(group_str, key=len)) + 2 |
|
|
group_str = "".join([f"{g:>{max_g_size}}" for g in group_str]) |
|
|
formatted_info += f"\n{group + ':':<{max_group_size}}{group_str}" |
|
|
return formatted_info |
|
|
|
|
|
def log_iter(self, epoch_num, iter_num, num_iters, iter_info, event="epoch"): |
|
|
output_info = f"{event.upper()} {epoch_num}, ITER {iter_num}/{num_iters}:" |
|
|
output_info += self.format_info(iter_info) |
|
|
self.logger.info(output_info) |
|
|
|
|
|
def log_epoch(self, epoch_info, epoch_num): |
|
|
output_info = f"EPOCH {epoch_num}:" |
|
|
output_info += self.format_info(epoch_info) |
|
|
self.logger.info(output_info) |
|
|
|
|
|
|
|
|
class WandbLogger: |
|
|
def __init__(self, config): |
|
|
wandb.login(key=os.environ['WANDB_KEY'].strip(), relogin=True) |
|
|
if config.train.resume_path == "": |
|
|
config_for_logger = omegaconf.OmegaConf.to_container(config) |
|
|
self.wandb_args = { |
|
|
"id": wandb.util.generate_id(), |
|
|
"project": config.exp.wandb_project, |
|
|
"name": config.exp.name, |
|
|
"config": config_for_logger, |
|
|
} |
|
|
wandb.init(**self.wandb_args, resume="allow") |
|
|
|
|
|
run_dir = wandb.run.dir |
|
|
print("run_dir", run_dir) |
|
|
|
|
|
code = wandb.Artifact("project-source", type="code") |
|
|
for path in glob.glob("**/*.py", recursive=True): |
|
|
if not path.startswith("wandb"): |
|
|
if os.path.basename(path) != path: |
|
|
code.add_dir( |
|
|
os.path.dirname(path), name=os.path.dirname(path) |
|
|
) |
|
|
else: |
|
|
code.add_file(os.path.basename(path), name=path) |
|
|
wandb.run.log_artifact(code) |
|
|
else: |
|
|
print(f"Resume training from {config.train.resume_path}") |
|
|
with open(config.train.resume_path, "r") as f: |
|
|
options = json.load(f) |
|
|
|
|
|
self.wandb_args = { |
|
|
"id": options['id'], |
|
|
"project": options['project'], |
|
|
"name": options['name'], |
|
|
"config": options['config'], |
|
|
} |
|
|
wandb.init(resume=True, **self.wandb_args) |
|
|
|
|
|
@staticmethod |
|
|
def log_epoch(iter_info, step): |
|
|
wandb.log( |
|
|
data={k: v.mean for k, v in iter_info.items()}, |
|
|
step=step + 1, |
|
|
commit=True, |
|
|
) |
|
|
|
|
|
@staticmethod |
|
|
def log_special_pics(pics, captions, paths): |
|
|
to_log = {} |
|
|
for i, path in enumerate(paths): |
|
|
to_log[path] = wandb.Image(pics[i], caption=captions[path]) |
|
|
wandb.log(to_log) |
|
|
|
|
|
|
|
|
class BlankWandbLogger: |
|
|
def __init__(self): |
|
|
self.wandb_args = None |
|
|
|
|
|
def log_epoch(*args, **kwars): |
|
|
pass |
|
|
|
|
|
def log_special_pics(*args, **kwars): |
|
|
pass |
|
|
|
|
|
|
|
|
class TrainigLogger: |
|
|
def __init__(self, config): |
|
|
self.console_logger = ConsoleLogger("") |
|
|
if config.exp.wandb == True: |
|
|
self.wandb_logger = WandbLogger(config) |
|
|
else: |
|
|
self.wandb_logger = BlankWandbLogger() |
|
|
|
|
|
self.trainig_steps = config.train.steps |
|
|
self.val_step = config.train.val_step |
|
|
|
|
|
def log_train_time_left(self, iter_info, step): |
|
|
float_iter_time = iter_info["duration/iter_train"].mean |
|
|
float_val_time = iter_info["duration/iter_val"].mean |
|
|
time_left = str( |
|
|
datetime.datetime.fromtimestamp( |
|
|
float_iter_time * (self.trainig_steps - step) |
|
|
+ float_val_time |
|
|
* ( |
|
|
(self.trainig_steps - step) // self.val_step |
|
|
) |
|
|
) |
|
|
- datetime.datetime.fromtimestamp(0) |
|
|
) |
|
|
|
|
|
print() |
|
|
print(f"Step {step}/{self.trainig_steps}") |
|
|
print(f"Time left: {time_left}") |
|
|
print(f"Time per step: {iter_info['duration/iter_train'].mean :.3f}") |
|
|
print() |
|
|
print() |
|
|
|
|
|
def save_train_logs(self, iter_info, step): |
|
|
self.wandb_logger.log_epoch(iter_info, step) |
|
|
self.console_logger.log_epoch(iter_info, step) |
|
|
|
|
|
self.log_train_time_left(iter_info, step) |
|
|
|
|
|
def save_validation_logs(self, orig_pics, method_pics, captions, special_paths): |
|
|
log_pics = [] |
|
|
for real_img, fake_img in zip(orig_pics, method_pics): |
|
|
concat_img = Image.new( |
|
|
"RGB", (real_img.width + fake_img.width, real_img.height) |
|
|
) |
|
|
concat_img.paste(real_img, (0, 0)) |
|
|
concat_img.paste(fake_img, (real_img.width, 0)) |
|
|
log_pics.append(concat_img) |
|
|
|
|
|
self.wandb_logger.log_special_pics(log_pics, captions, special_paths) |
|
|
|
