231 lines
9.0 KiB
Python
231 lines
9.0 KiB
Python
# -*- coding: utf-8 -*-
|
|
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
import json
|
|
import os
|
|
|
|
import numpy as np
|
|
import torch
|
|
from tqdm import tqdm
|
|
|
|
from scepter.modules.data.dataset import DATASETS
|
|
from scepter.modules.solver.hooks.hook import Hook
|
|
from scepter.modules.solver.hooks.registry import HOOKS
|
|
from scepter.modules.utils.config import dict_to_yaml
|
|
from scepter.modules.utils.distribute import barrier, gather_data, we
|
|
from scepter.modules.utils.file_system import FS
|
|
from scepter.modules.utils.math_plot import plot_multi_curves
|
|
|
|
_DEFAULT_VAL_PRIORITY = 200
|
|
|
|
|
|
def float_format(o):
|
|
if isinstance(o, float):
|
|
return f"{o: .6f}"
|
|
raise TypeError(f"Type {type(o)} not serializable")
|
|
|
|
|
|
@HOOKS.register_class()
|
|
class ValLossHook(Hook):
|
|
para_dict = [{
|
|
'PRIORITY': {
|
|
'value': _DEFAULT_VAL_PRIORITY,
|
|
'description': 'The priority for processing!'
|
|
},
|
|
'VAL_INTERVAL': {
|
|
'value': 1000,
|
|
'description': 'the interval for log print!'
|
|
},
|
|
'VAL_LIMITATION_SIZE': {
|
|
'value': 1000000,
|
|
'description': 'the limitation size for validation!'
|
|
},
|
|
'VAL_SEED': {
|
|
'value': 2025,
|
|
'description': 'the validation seed for t or generator sample!'
|
|
}
|
|
}]
|
|
|
|
def __init__(self, cfg, logger=None):
|
|
super(ValLossHook, self).__init__(cfg, logger=logger)
|
|
self.priority = cfg.get('PRIORITY', _DEFAULT_VAL_PRIORITY)
|
|
self.val_interval = cfg.get('VAL_INTERVAL', 1000)
|
|
self.val_dim = cfg.get('VAL_DIM', 'all')
|
|
self.meta_field = cfg.get('META_FIELD', ['edit_type', 'data_type'])
|
|
self.save_folder = cfg.get('SAVE_FOLDER', 'val_loss')
|
|
self.val_limitation_size = cfg.get('VAL_LIMITATION_SIZE', 1000000)
|
|
self.val_seed = cfg.get('VAL_SEED', 2025)
|
|
self.data = DATASETS.build(cfg.DATA, logger=logger)
|
|
|
|
def before_all_iter(self, solver):
|
|
solver.eval_mode()
|
|
self.eval_set_size = len(self.data.dataset)
|
|
if self.eval_set_size > self.val_limitation_size:
|
|
self.logger.info(
|
|
f"The samples number {self.eval_set_size} of validation set "
|
|
f"should not great than {self.val_limitation_size}")
|
|
assert self.eval_set_size < self.val_limitation_size
|
|
if not hasattr(solver, 'run_step_val'):
|
|
self.logger.info(
|
|
f"The val-loss hook should have the function run_step_val" # noqa
|
|
) # noqa
|
|
assert hasattr(solver, 'run_step_val')
|
|
if not self.data.batch_size == 1:
|
|
self.logger.info(
|
|
f"The batch_size of validation set should be 1 " # noqa
|
|
f"when you use the validation hook to make the results deterministic." # noqa
|
|
)
|
|
assert self.data.batch_size == 1
|
|
timestamp_generator = torch.Generator(device=we.device_id)
|
|
timestamp_generator.manual_seed(self.val_seed)
|
|
u = torch.rand((self.eval_set_size, ),
|
|
device=we.device_id,
|
|
generator=timestamp_generator)
|
|
self.t = (u * (solver.timesteps - 1)).round().long()
|
|
solver.val_interval = self.val_interval
|
|
solver.train_mode()
|
|
|
|
def get_val_loss(self, solver, step):
|
|
all_loss = []
|
|
# batch-size must be 1
|
|
for batch_data in tqdm(self.data.dataloader):
|
|
# generate t list
|
|
sample_id = int(batch_data['sample_id'][0])
|
|
meta_info = {m_f: batch_data[m_f][0] for m_f in self.meta_field}
|
|
meta_info['sample_id'] = sample_id
|
|
|
|
batch_data['t'] = torch.stack(
|
|
[self.t[sample_id % self.eval_set_size]])
|
|
noise_generator = torch.Generator(device=we.device_id)
|
|
noise_generator.manual_seed(sample_id + 10000 * self.val_seed)
|
|
# get generator according to the sample_id
|
|
with torch.no_grad():
|
|
loss = solver.run_step_val(batch_data, noise_generator)
|
|
meta_info['loss'] = float(loss[sample_id])
|
|
all_loss.append(meta_info)
|
|
all_loss = json.dumps(all_loss, default=float_format)
|
|
all_loss = gather_data([all_loss])
|
|
if we.rank == 0:
|
|
reduce_loss = []
|
|
for loss in all_loss:
|
|
reduce_loss.extend(json.loads(loss))
|
|
compute_results = self.compute_avg_loss(reduce_loss)
|
|
self.save_record(solver, compute_results, reduce_loss, step)
|
|
return
|
|
|
|
def compute_avg_loss(self, loss_list):
|
|
all_avg_ls = []
|
|
avg_ls = {}
|
|
for ls in loss_list:
|
|
for m_f in self.meta_field:
|
|
m_f_v = ls[m_f]
|
|
ls_key = m_f + '_' + m_f_v
|
|
if ls_key not in avg_ls:
|
|
avg_ls[ls_key] = []
|
|
avg_ls[ls_key].append(ls['loss'])
|
|
all_avg_ls.append(ls['loss'])
|
|
compute_results = {
|
|
'all': sum(all_avg_ls) / len(all_avg_ls),
|
|
}
|
|
compute_results.update(
|
|
{m_f: sum(avg_ls[m_f]) / len(avg_ls[m_f])
|
|
for m_f in avg_ls})
|
|
return compute_results
|
|
|
|
def save_record(self, solver, compute_results, all_loss, step):
|
|
save_folder = os.path.join(solver.work_dir, self.save_folder)
|
|
# save history
|
|
save_history = os.path.join(save_folder, 'history.json')
|
|
|
|
draw_curve = False
|
|
|
|
if FS.exists(save_history):
|
|
results = json.loads(FS.get_object(save_history).decode())
|
|
all_loss = {loss['sample_id']: loss for loss in all_loss}
|
|
for loss in results['detail']:
|
|
loss['loss'] = {int(k): v for k, v in loss['loss'].items()}
|
|
loss['loss'][step] = all_loss[loss['sample_id']]['loss']
|
|
for k, v in compute_results.items():
|
|
results['summary'][k] = {
|
|
int(kk): vv
|
|
for kk, vv in results['summary'][k].items()
|
|
}
|
|
results['summary'][k][step] = v
|
|
draw_curve = True
|
|
else:
|
|
results = {'detail': [], 'summary': {}}
|
|
for loss in all_loss:
|
|
loss_v = loss.pop('loss')
|
|
loss['loss'] = {step: loss_v}
|
|
results['detail'].append(loss)
|
|
for k, v in compute_results.items():
|
|
if k not in results['summary']:
|
|
results['summary'][k] = {}
|
|
results['summary'][k][step] = v
|
|
#
|
|
FS.put_object(
|
|
json.dumps(results, default=float_format).encode(), save_history)
|
|
# plot current curve
|
|
if draw_curve:
|
|
self.plot_results(results['summary'],
|
|
os.path.join(save_folder, 'curve'))
|
|
# print current log
|
|
print_msg = ''
|
|
for k, v in compute_results.items():
|
|
print_msg += f"{k}: {v: .4f} "
|
|
self.logger.info(f"Step {step} validation loss: {print_msg}")
|
|
|
|
def plot_results(self, plot_data, save_folder):
|
|
y = []
|
|
steps = []
|
|
# one image
|
|
for label, curve_data in plot_data.items():
|
|
curve_data = [[step, value] for step, value in curve_data.items()]
|
|
curve_data.sort(key=lambda x: x[0])
|
|
steps = [step for step, value in curve_data]
|
|
value = [value for step, value in curve_data]
|
|
k_y = [{'data': np.array(value), 'label': label}]
|
|
save_path = os.path.join(save_folder, 'detail', f"{label}.png")
|
|
with FS.put_to(save_path) as local_file:
|
|
plot_multi_curves(x=np.array(steps),
|
|
y=k_y,
|
|
x_label='steps',
|
|
y_label=None,
|
|
title=f"{label}'s validation loss",
|
|
save_path=local_file)
|
|
y = y + k_y
|
|
if len(steps) > 0:
|
|
save_path = os.path.join(save_folder, f"summary.png") # noqa
|
|
with FS.put_to(save_path) as local_file:
|
|
plot_multi_curves(
|
|
x=np.array(steps),
|
|
y=y,
|
|
x_label='steps',
|
|
y_label=None,
|
|
title=f"validation loss", # noqa
|
|
save_path=local_file)
|
|
|
|
def after_iter(self, solver):
|
|
if solver.mode == 'train' and solver.total_iter % self.val_interval == 0:
|
|
step = solver.total_iter
|
|
solver.eval_mode()
|
|
self.get_val_loss(solver, step)
|
|
solver.train_mode()
|
|
torch.cuda.synchronize()
|
|
barrier()
|
|
|
|
def after_all_iter(self, solver):
|
|
if solver.mode == 'train':
|
|
step = solver.total_iter
|
|
solver.eval_mode()
|
|
self.get_val_loss(solver, step)
|
|
solver.train_mode()
|
|
torch.cuda.synchronize()
|
|
barrier()
|
|
|
|
@staticmethod
|
|
def get_config_template():
|
|
return dict_to_yaml('HOOK',
|
|
__class__.__name__,
|
|
ValLossHook.para_dict,
|
|
set_name=True)
|