def __init__(self, grad_clip=None):
super().__init__()
self._grad_clip = grad_clip
self._level = PriorityStatus.HIGH
self._user_scale = False if self._grad_clip is None else getattr(
self._grad_clip, 'use_scale', False)
self._scaler = GradScaler(enabled=False) if self._user_scale else None
def load_model(self, checkpoint=None):
print('Loading model from', self.config['model'])
self.model, self.tokenizer = instantiate_model_and_tokenizer(
self.config['model'],
additional_tokens_smart_init=self.config['smart_init'],
dropout=self.config['dropout'],
attention_dropout=self.config['attention_dropout'],
penman_linearization=self.config['penman_linearization'],
collapse_name_ops=self.config['collapse_name_ops'],
use_pointer_tokens=self.config['use_pointer_tokens'],
raw_graph=self.config['raw_graph'])
self.model.to(self.device)
# Load optimization components
self.optimizer = AdamW(self.model.parameters(),
lr=self.config['learning_rate'],
weight_decay=self.config['weight_decay'])
self.scheduler = transformers.get_constant_schedule_with_warmup(
self.optimizer, num_warmup_steps=self.config['warmup_steps'])
self.scaler = GradScaler(enabled=self.config['fp16'])
# Reload checkpoint model weights and optimizer params if loading from a checkpoint
if checkpoint is not None:
print('Checkpoint %s restored' % checkpoint)
load_state_dict_from_checkpoint(checkpoint, self.model,
self.optimizer, self.scheduler)
# Try to load the smatch score and last_epoch from the config in the model directory.
try:
with open(os.path.join(self.model_dir, 'config.json')) as f:
model_config = json.load(f)
self.best_smatch = model_config['smatch_dev']
self.start_epoch = model_config['last_epoch'] + 1
except:
logger.exception(
'Unable to load config file in model directory')
def __init__(self, cfg):
self.cfg = cfg
self.paths = cfg['paths']
self.net_params = cfg['net']
self.train_params = cfg['train']
self.trans_params = cfg['train']['transforms']
self.checkpoints = self.paths['checkpoints']
Path(self.checkpoints).mkdir(parents=True, exist_ok=True)
shutil.copyfile('config.yaml', f'{self.checkpoints}/config.yaml')
self.update_interval = self.paths['update_interval']
# amp training
self.use_amp = self.train_params['mixed_precision']
self.scaler = GradScaler() if self.use_amp else None
# data setup
dataset_name = self.train_params['dataset']
self.use_multi = dataset_name == 'multi'
print(f'Using dataset: {dataset_name}')
self.train_dataset = get_pedestrian_dataset(
dataset_name,
self.paths,
augment=get_train_transforms(self.trans_params),
mode='train',
multi_datasets=self.train_params['multi_datasets']
if self.use_multi else None)
print(f'Train dataset: {len(self.train_dataset)} samples')
self.val_dataset = get_pedestrian_dataset(
dataset_name,
self.paths,
augment=get_val_transforms(self.trans_params),
mode='val',
multi_datasets=self.train_params['multi_datasets']
if self.use_multi else None)
print(f'Val dataset: {len(self.val_dataset)} samples')
tests_data = self.train_params['test_datasets']
self.test_datasets = [
get_pedestrian_dataset(d_name,
self.paths,
augment=get_test_transforms(
self.trans_params),
mode='test') for d_name in tests_data
]
self.criterion = AnchorFreeLoss(self.train_params)
self.writer = Writer(self.paths['log_dir'])
print('Tensorboard logs are saved to: {}'.format(
self.paths['log_dir']))
self.sched_type = self.train_params['scheduler']
self.scheduler = None
self.optimizer = None
def __init__(
self,
enabled: bool = False,
max_norm: Optional[float] = None,
) -> None:
self.grad_scaler = GradScaler(enabled=enabled)
self.enabled = enabled
self.max_norm = max_norm
_logger.info("amp: %s", self.enabled)
if self.max_norm:
_logger.info(
"you are using grad clip, don't forget to pass params in")
class OptimizerHook(HookBase):
def __init__(self, grad_clip=None):
super().__init__()
self._grad_clip = grad_clip
self._level = PriorityStatus.HIGH
self._user_scale = False if self._grad_clip is None else getattr(
self._grad_clip, 'use_scale', False)
self._scaler = GradScaler(enabled=False) if self._user_scale else None
def _clip_grad_norm(self) -> None:
clip_norm_params = list(
filter(lambda parm: parm.requires_grad and parm.grad is not None,
self.trainer.model.parameters()))
if len(clip_norm_params) == 0:
return
else:
if hasattr(self._grad_clip, 'clip_norm_mode'):
scale = self._scaler.get_scale() if self._user_scale else 1.0
max_norm = self._grad_clip.max_grad_l2_norm * scale
grad_norm = clip_grad.clip_grad_norm(clip_norm_params,
max_norm)
else:
grad_norm = clip_grad.clip_grad_norm_(clip_norm_params,
**self._grad_clip)
self.trainer.log_buffer.put_scalar('grad_norm', float(grad_norm))
def after_train_iter(self):
self.trainer.output['loss'] /= self.trainer.gradient_accumulation_steps
self.trainer.output['loss'].backward()
if self._grad_clip is not None:
self._clip_grad_norm()
if (self.trainer.iter +
1) % self.trainer.gradient_accumulation_steps == 0:
if self._user_scale:
self._scaler.step(self.trainer.optimizer)
self._scaler.update()
else:
self.trainer.optimizer.step()
def before_train_iter(self):
if self.trainer.iter == 0:
is_clean = True
elif self.trainer.iter % self.trainer.gradient_accumulation_steps == 0:
is_clean = True
else:
is_clean = False
if is_clean:
self.trainer.optimizer.zero_grad()
def fnTrain(
loader: DataLoader,
device: str,
model: nn.Module,
optimizer: Optimizer,
fnLoss,
scaler: GradScaler,
) -> float:
runningLoss = 0
for _, (data, targets) in enumerate(loader):
data = data.to(device=device)
targets = targets.float().unsqueeze(1).to(device=device)
with torch.cuda.amp.autocast():
predictions = model(data)
loss = fnLoss(predictions, targets)
optimizer.zero_grad()
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
# print(f"batch {idxBatch+ 1} loss {loss.item()}")
runningLoss += loss.item()
return runningLoss / len(loader)
def train_on_epoch(self, fold, epoch):
logging.info(f'Fold-[{fold}] ==> Training on epoch{epoch}...')
trn_loss_list = []
trn_acc_list = []
trn_f1_list = []
self.model.train()
if self.cfg.mixed_precision:
scaler = GradScaler()
else:
scaler = None
self.reset_scoring_dict()
with tqdm(self.trn_dl, total=len(self.trn_dl), unit="batch") as train_bar:
for batch, sample in enumerate(train_bar):
train_bar.set_description(f"Fold-[{fold}|{self.cfg.n_fold}] ==> Train Epoch [{str(epoch).zfill(len(str(self.cfg.t_epoch)))}|{self.cfg.t_epoch}]")
result = self.step(sample, scaler)
batch_f1 = result['batch_f1']
batch_acc = result['batch_acc']
loss = result['loss']
if not torch.isfinite(loss):
print(loss, sample, result['logit'], sample['image'].shape, sample['label'].shape)
raise ValueError('WARNING: non-finite loss, ending training ')
trn_f1_list.append(batch_f1)
trn_acc_list.append(batch_acc)
trn_loss_list.append(loss.item())
trn_f1 = np.mean(trn_f1_list)
trn_acc = np.mean(trn_acc_list)
trn_loss = np.mean(trn_loss_list)
if batch % self.cfg.log_interval == 0 or batch == len(self.trn_dl)-1:
logging.info(f"Fold-[{fold}] ==> <Train> Epoch: [{str(epoch).zfill(len(str(self.cfg.t_epoch)))}|{str(self.cfg.t_epoch)}] Batch: [{str(batch).zfill(len(str(len(self.trn_dl))))}|{len(self.trn_dl)}]\t Train Acc: {trn_acc}\t Train F1: {trn_f1}\t Train Loss: {trn_loss}")
train_bar.set_postfix(train_loss=trn_loss, train_acc=trn_acc, train_f1=trn_f1)
if self.cfg.sched_type == "onecycle":
self.sched.step()
if self.cfg.sched_type == "cosine": self.sched.step()
reports = report(self.scoring_dict["preds"], self.scoring_dict["labels"])
logging.info(f"Fold-[{fold}] ==> <Train> Epoch: [{str(epoch).zfill(len(str(self.cfg.t_epoch)))}|{str(self.cfg.t_epoch)}] REPOST\n{reports}\n")
class Fp16OptimizerHook(OptimizerHook):
def __init__(self, grad_clip=None, grad_scaler_config=None):
super().__init__(grad_clip)
self._grad_scaler_config = grad_scaler_config
self._scaler = None
def before_train(self):
if self._grad_scaler_config is None:
self._scaler = GradScaler()
else:
self._scaler = GradScaler(**self._grad_scaler_config)
def after_train_iter(self):
loss = self.trainer.output[
'loss'] / self.trainer.gradient_accumulation_steps
self._scaler.scale(loss).backward()
if self._grad_clip is not None:
self._scaler.unscale_(self.trainer.optimizer)
self._clip_grad_norm()
if (self.trainer.iter +
1) % self.trainer.gradient_accumulation_steps == 0:
self._scaler.step(self.trainer.optimizer)
self._scaler.update()
def valid_on_epoch(self, fold, epoch):
logging.info(f'Fold-[{fold}] ==> Validation on epoch{epoch}...')
val_loss_list = []
val_acc_list = []
val_f1_list = []
self.model.eval()
if self.cfg.mixed_precision: scaler = GradScaler()
else: scaler = None
self.reset_scoring_dict()
with torch.no_grad():
with tqdm(self.val_dl, total=len(self.val_dl), unit="batch") as valid_bar:
for batch, sample in enumerate(valid_bar):
valid_bar.set_description(f"Fold-[{fold}|{self.cfg.n_fold}] ==> Valid Epoch [{str(epoch).zfill(len(str(self.cfg.t_epoch)))}|{self.cfg.t_epoch}]")
result = self.step(sample, scaler, valid=True)
batch_f1 = result['batch_f1']
batch_acc = result['batch_acc']
loss = result['loss']
val_f1_list.append(batch_f1)
val_acc_list.append(batch_acc)
val_loss_list.append(loss.item())
val_f1 = np.mean(val_f1_list)
val_acc = np.mean(val_acc_list)
val_loss = np.mean(val_loss_list)
if batch % self.cfg.log_interval == 0 or batch == len(self.val_dl)-1:
logging.info(f"Fold-[{fold}] ==> <Valid> Epoch: [{str(epoch).zfill(len(str(self.cfg.t_epoch)))}|{str(self.cfg.t_epoch)}] Batch: [{str(batch).zfill(len(str(len(self.val_dl))))}|{len(self.val_dl)}]\t Valid Acc: {val_acc}\t Valid F1: {val_f1}\t Valid Loss: {val_loss}")
valid_bar.set_postfix(valid_loss=val_loss,valid_acc=val_acc, valid_f1=val_f1)
if self.cfg.sched_type == "plateau": self.sched.step(val_loss)
reports = report(self.scoring_dict["preds"], self.scoring_dict["labels"])
logging.info(f"Fold-[{fold}] ==> <Valid> Epoch: [{str(epoch).zfill(len(str(self.cfg.t_epoch)))}|{str(self.cfg.t_epoch)}] REPOST\n{reports}\n")
return val_loss, val_acc, val_f1
def __init__(self, cfg):
super().__init__()
logger = logging.getLogger("fastreid")
if not logger.isEnabledFor(
logging.INFO): # if setup_logger is not called for fastreid
setup_logger()
# Create datasets
logger.info("==> Load source-domain dataset")
self.src = src = self.load_dataset(cfg.DATASETS.SRC)
self.src_pid_nums = src.get_num_pids(src.train)
logger.info("==> Load target-domain dataset")
self.tgt = tgt = self.load_dataset(cfg.DATASETS.TGT)
self.tgt_nums = len(tgt.train)
# Create model
self.model = self.build_model(cfg,
load_model=False,
show_model=False,
use_dsbn=True)
# Create hybrid memorys
self.hm = HybridMemory(num_features=cfg.MODEL.BACKBONE.FEAT_DIM,
num_samples=self.src_pid_nums + self.tgt_nums,
temp=cfg.MEMORY.TEMP,
momentum=cfg.MEMORY.MOMENTUM,
use_half=cfg.SOLVER.AMP.ENABLED).cuda()
# Initialize source-domain class centroids
logger.info(
"==> Initialize source-domain class centroids in the hybrid memory"
)
with inference_context(self.model), torch.no_grad():
src_train = self.build_dataset(cfg,
src.train,
is_train=False,
relabel=False,
with_mem_idx=False)
src_init_feat_loader = self.build_test_loader(cfg, src_train)
src_fname_feat_dict, _ = extract_features(self.model,
src_init_feat_loader)
src_feat_dict = collections.defaultdict(list)
for f, pid, _ in sorted(src.train):
src_feat_dict[pid].append(src_fname_feat_dict[f].unsqueeze(0))
src_centers = [
torch.cat(src_feat_dict[pid], 0).mean(0)
for pid in sorted(src_feat_dict.keys())
]
src_centers = torch.stack(src_centers, 0)
src_centers = F.normalize(src_centers, dim=1)
# Initialize target-domain instance features
logger.info(
"==> Initialize target-domain instance features in the hybrid memory"
)
with inference_context(self.model), torch.no_grad():
tgt_train = self.build_dataset(cfg,
tgt.train,
is_train=False,
relabel=False,
with_mem_idx=False)
tgt_init_feat_loader = self.build_test_loader(cfg, tgt_train)
tgt_fname_feat_dict, _ = extract_features(self.model,
tgt_init_feat_loader)
tgt_features = torch.cat([
tgt_fname_feat_dict[f].unsqueeze(0)
for f, _, _ in sorted(self.tgt.train)
], 0)
tgt_features = F.normalize(tgt_features, dim=1)
self.hm.features = torch.cat((src_centers, tgt_features), dim=0).cuda()
del (src_train, src_init_feat_loader, src_fname_feat_dict,
src_feat_dict, src_centers, tgt_train, tgt_init_feat_loader,
tgt_fname_feat_dict, tgt_features)
# Optimizer
self.optimizer, self.param_wrapper = self.build_optimizer(
cfg, self.model)
# For training, wrap with DDP. But don't need this for inference.
if comm.get_world_size() > 1:
# ref to https://github.com/pytorch/pytorch/issues/22049 to set `find_unused_parameters=True`
# for part of the parameters is not updated.
self.model = DistributedDataParallel(
self.model,
device_ids=[comm.get_local_rank()],
broadcast_buffers=False,
find_unused_parameters=True)
# Learning rate scheduler
self.iters_per_epoch = cfg.SOLVER.ITERS
self.scheduler = self.build_lr_scheduler(cfg, self.optimizer,
self.iters_per_epoch)
# Assume no other objects need to be checkpointed.
# We can later make it checkpoint the stateful hooks
self.checkpointer = Checkpointer(
# Assume you want to save checkpoints together with logs/statistics
self.model,
cfg.OUTPUT_DIR,
save_to_disk=comm.is_main_process(),
optimizer=self.optimizer,
**self.scheduler,
)
self.start_epoch = 0
self.max_epoch = cfg.SOLVER.MAX_EPOCH
self.max_iter = self.max_epoch * self.iters_per_epoch
self.warmup_iters = cfg.SOLVER.WARMUP_ITERS
self.delay_epochs = cfg.SOLVER.DELAY_EPOCHS
self.cfg = cfg
self.register_hooks(self.build_hooks())
if cfg.SOLVER.AMP.ENABLED:
unsupported = "AMPTrainer does not support single-process multi-device training!"
if isinstance(self.model, DistributedDataParallel):
assert not (self.model.device_ids
and len(self.model.device_ids) > 1), unsupported
from torch.cuda.amp.grad_scaler import GradScaler
self.grad_scaler = GradScaler()
else:
self.grad_scaler = None
def __init__(self, cfg):
super().__init__()
logger = logging.getLogger("fastreid")
if not logger.isEnabledFor(
logging.INFO): # if setup_logger is not called for fastreid
setup_logger()
logger.info("==> Load target-domain dataset")
self.tgt = tgt = self.load_dataset(cfg.DATASETS.TGT)
self.tgt_nums = len(tgt.train)
cfg = self.auto_scale_hyperparams(cfg, self.tgt_nums)
# Create model
self.model = self.build_model(cfg,
load_model=cfg.MODEL.PRETRAIN,
show_model=True,
use_dsbn=False)
# Optimizer
self.optimizer, self.param_wrapper = self.build_optimizer(
cfg, self.model)
# For training, wrap with DDP. But don't need this for inference.
if comm.get_world_size() > 1:
# ref to https://github.com/pytorch/pytorch/issues/22049 to set `find_unused_parameters=True`
# for part of the parameters is not updated.
self.model = DistributedDataParallel(
self.model,
device_ids=[comm.get_local_rank()],
broadcast_buffers=False,
find_unused_parameters=True)
# Learning rate scheduler
self.iters_per_epoch = cfg.SOLVER.ITERS
self.scheduler = self.build_lr_scheduler(cfg, self.optimizer,
self.iters_per_epoch)
# Assume no other objects need to be checkpointed.
# We can later make it checkpoint the stateful hooks
self.checkpointer = Checkpointer(
# Assume you want to save checkpoints together with logs/statistics
self.model,
cfg.OUTPUT_DIR,
save_to_disk=comm.is_main_process(),
optimizer=self.optimizer,
**self.scheduler,
)
self.start_epoch = 0
self.max_epoch = cfg.SOLVER.MAX_EPOCH
self.max_iter = self.max_epoch * self.iters_per_epoch
self.warmup_iters = cfg.SOLVER.WARMUP_ITERS
self.delay_epochs = cfg.SOLVER.DELAY_EPOCHS
self.cfg = cfg
self.register_hooks(self.build_hooks())
if cfg.SOLVER.AMP.ENABLED:
unsupported = "AMPTrainer does not support single-process multi-device training!"
if isinstance(self.model, DistributedDataParallel):
assert not (self.model.device_ids
and len(self.model.device_ids) > 1), unsupported
from torch.cuda.amp.grad_scaler import GradScaler
self.grad_scaler = GradScaler()
else:
self.grad_scaler = None
def train(rank, cfg: TrainConfig):
if cfg.distributed.n_gpus_per_node > 1:
init_process_group(backend=cfg.distributed.dist_backend,
init_method=cfg.distributed.dist_url,
world_size=cfg.distributed.n_nodes *
cfg.distributed.n_gpus_per_node,
rank=rank)
device = torch.device(f'cuda:{rank:d}')
model = ConvRNNEmbedder(cfg.model_cfg).to(device)
loss_fn = GE2ELoss(device).to(device)
logging.info(f"Initialized rank {rank}")
if rank == 0:
logging.getLogger().setLevel(logging.INFO)
logging.info(f"Model initialized as:\n {model}")
os.makedirs(cfg.checkpoint_path, exist_ok=True)
logging.info(f"checkpoints directory : {cfg.checkpoint_path}")
logging.info(
f"Model has {sum([p.numel() for p in model.parameters()]):,d} parameters."
)
steps = 0
if cfg.resume_checkpoint != '' and os.path.isfile(cfg.resume_checkpoint):
state_dict = torch.load(cfg.resume_checkpoint, map_location=device)
model.load_state_dict(state_dict['model_state_dict'])
loss_fn.load_state_dict(state_dict['loss_fn_state_dict'])
steps = state_dict['steps'] + 1
last_epoch = state_dict['epoch']
print(
f"Checkpoint loaded from {cfg.resume_checkpoint}. Resuming training from {steps} steps at epoch {last_epoch}"
)
else:
state_dict = None
last_epoch = -1
if cfg.distributed.n_gpus_per_node * cfg.distributed.n_nodes > 1:
if rank == 0: logging.info("Multi-gpu detected")
model = DDP(model, device_ids=[rank]).to(device)
loss_fn = DDP(loss_fn, device_ids=[rank]).to(device)
optim = torch.optim.AdamW(chain(model.parameters(), loss_fn.parameters()),
1.0,
betas=cfg.betas)
if state_dict is not None:
optim.load_state_dict(state_dict['optim_state_dict'])
train_df, valid_df = pd.read_csv(cfg.train_csv), pd.read_csv(cfg.valid_csv)
trainset = UtteranceDS(train_df, cfg.sample_rate, cfg.n_uttr_per_spk)
train_sampler = DistributedSampler(
trainset
) if cfg.distributed.n_gpus_per_node * cfg.distributed.n_nodes > 1 else None
train_loader = DataLoader(trainset,
num_workers=cfg.num_workers,
shuffle=False,
sampler=train_sampler,
batch_size=cfg.batch_size,
pin_memory=False,
drop_last=True,
collate_fn=SpecialCollater(
cfg.min_seq_len, cfg.max_seq_len))
if rank == 0:
validset = UtteranceDS(valid_df, cfg.sample_rate, cfg.n_uttr_per_spk)
validation_loader = DataLoader(validset,
num_workers=cfg.num_workers,
shuffle=False,
sampler=None,
batch_size=cfg.batch_size,
pin_memory=False,
drop_last=True,
collate_fn=SpecialCollater(
cfg.min_seq_len, cfg.max_seq_len))
sw = SummaryWriter(os.path.join(cfg.checkpoint_path, 'logs'))
total_iters = cfg.n_epochs * len(train_loader)
def sched_lam(x):
return lin_one_cycle(cfg.start_lr, cfg.max_lr, cfg.end_lr,
cfg.warmup_pct, total_iters, x)
scheduler = torch.optim.lr_scheduler.LambdaLR(optim,
lr_lambda=[sched_lam],
last_epoch=steps - 1)
if state_dict is not None:
scheduler.load_state_dict(state_dict['scheduler_state_dict'])
if cfg.fp16:
scaler = GradScaler()
if state_dict is not None and 'scaler_state_dict' in state_dict:
scaler.load_state_dict(state_dict['scaler_state_dict'])
model.train()
if rank == 0:
mb = master_bar(range(max(0, last_epoch), cfg.n_epochs))
smooth_loss = None
else:
mb = range(max(0, last_epoch), cfg.n_epochs)
for epoch in mb:
if rank == 0:
start = time.time()
mb.write("Epoch: {}".format(epoch + 1))
if cfg.distributed.n_gpus_per_node * cfg.distributed.n_nodes > 1:
train_sampler.set_epoch(epoch)
if rank == 0:
pb = progress_bar(enumerate(train_loader),
total=len(train_loader),
parent=mb)
else:
pb = enumerate(train_loader)
for i, batch in pb:
if rank == 0: start_b = time.time()
x, xlen = batch
x = x.to(device, non_blocking=True)
xlen = xlen.to(device, non_blocking=True)
optim.zero_grad()
with torch.cuda.amp.autocast(enabled=cfg.fp16):
embeds = model(x, xlen)
loss = loss_fn(embeds)
if cfg.fp16:
scaler.scale(loss).backward()
scaler.unscale_(optim)
gnorm = torch.nn.utils.clip_grad.clip_grad_norm_(
model.parameters(), cfg.grad_clip)
torch.nn.utils.clip_grad.clip_grad_norm_(
loss_fn.parameters(), cfg.grad_clip / 2)
scaler.step(optim)
scaler.update()
else:
loss.backward()
gnorm = torch.nn.utils.clip_grad.clip_grad_norm_(
model.parameters(), cfg.grad_clip)
torch.nn.utils.clip_grad.clip_grad_norm_(
loss_fn.parameters(), cfg.grad_clip / 2)
optim.step()
if rank == 0:
if smooth_loss is None: smooth_loss = float(loss.item())
else:
smooth_loss = smooth_loss + 0.1 * (float(loss.item()) -
smooth_loss)
# STDOUT logging
if steps % cfg.stdout_interval == 0:
mb.write('steps : {:,d}, loss : {:4.3f}, sec/batch : {:4.3f}, peak mem: {:5.2f}GB'. \
format(steps, loss.item(), time.time() - start_b, torch.cuda.max_memory_allocated()/1e9))
mb.child.comment = 'steps : {:,d}, loss : {:4.3f}, sec/batch : {:4.3f}'. \
format(steps, loss.item(), time.time() - start_b)
# mb.write(f"lr = {float(optim.param_groups[0]['lr'])}")
# checkpointing
if steps % cfg.checkpoint_interval == 0 and steps != 0:
checkpoint_path = f"{cfg.checkpoint_path}/ckpt_{steps:08d}.pt"
torch.save(
{
'model_state_dict':
(model.module if cfg.distributed.n_gpus_per_node *
cfg.distributed.n_nodes > 1 else
model).state_dict(),
'loss_fn_state_dict':
(loss_fn.module
if cfg.distributed.n_gpus_per_node *
cfg.distributed.n_nodes > 1 else
loss_fn).state_dict(),
'optim_state_dict':
optim.state_dict(),
'scheduler_state_dict':
scheduler.state_dict(),
'scaler_state_dict':
(scaler.state_dict() if cfg.fp16 else None),
'steps':
steps,
'epoch':
epoch
}, checkpoint_path)
logging.info(f"Saved checkpoint to {checkpoint_path}")
# Tensorboard summary logging
if steps % cfg.summary_interval == 0:
sw.add_scalar("training/loss_smooth", smooth_loss, steps)
sw.add_scalar("training/loss_raw", loss.item(), steps)
sw.add_scalar(
"ge2e/w",
float((loss_fn.module
if cfg.distributed.n_gpus_per_node *
cfg.distributed.n_nodes > 1 else
loss_fn).w.item()), steps)
sw.add_scalar(
"ge2e/b",
float((loss_fn.module
if cfg.distributed.n_gpus_per_node *
cfg.distributed.n_nodes > 1 else
loss_fn).b.item()), steps)
sw.add_scalar("opt/lr", float(optim.param_groups[0]['lr']),
steps)
sw.add_scalar('opt/grad_norm', float(gnorm), steps)
# Validation
if steps % cfg.validation_interval == 0 and steps != 0:
model.eval()
loss_fn.eval()
torch.cuda.empty_cache()
val_err_tot = 0
flat_embeds = []
flat_lbls = []
with torch.no_grad():
for j, batch in progress_bar(
enumerate(validation_loader),
total=len(validation_loader),
parent=mb):
x, xlen = batch
embeds = model(x.to(device), xlen.to(device))
val_err_tot += loss_fn(embeds)
if j <= 2:
lbls = [
f'spk-{j}-{indr:03d}'
for indr in range(cfg.batch_size)
for _ in range(cfg.n_uttr_per_spk)
]
fembeds = embeds.view(
cfg.batch_size * cfg.n_uttr_per_spk,
cfg.model_cfg.fc_dim)
flat_embeds.append(fembeds.cpu())
flat_lbls.extend(lbls)
elif j == 3:
flat_embeds = torch.cat(flat_embeds, dim=0)
sw.add_embedding(flat_embeds,
metadata=flat_lbls,
global_step=steps)
val_err = val_err_tot / (j + 1)
sw.add_scalar("validation/loss", val_err, steps)
mb.write(
f"validation run complete at {steps:,d} steps. validation loss: {val_err:5.4f}"
)
model.train()
loss_fn.train()
sw.add_scalar("memory/max_allocated_gb",
torch.cuda.max_memory_allocated() / 1e9,
steps)
sw.add_scalar("memory/max_reserved_gb",
torch.cuda.max_memory_reserved() / 1e9,
steps)
torch.cuda.reset_peak_memory_stats()
torch.cuda.reset_accumulated_memory_stats()
steps += 1
scheduler.step()
if rank == 0:
print('Time taken for epoch {} is {} sec\n'.format(
epoch + 1, int(time.time() - start)))
sw.add_hparams(flatten_cfg(cfg),
metric_dict={'validation/loss': val_err},
run_name=f'run-{cfg.checkpoint_path}')
print("Training completed!")
class Amp:
def __init__(
self,
enabled: bool = False,
max_norm: Optional[float] = None,
) -> None:
self.grad_scaler = GradScaler(enabled=enabled)
self.enabled = enabled
self.max_norm = max_norm
_logger.info("amp: %s", self.enabled)
if self.max_norm:
_logger.info(
"you are using grad clip, don't forget to pass params in")
def autocast(self):
return autocast(enabled=self.enabled)
def scale(self, outputs: TensorOrIterableTensors) -> TensorOrIterableTensors:
return self.grad_scaler.scale(outputs)
def unscale_(self, optimizer: Optimizer):
return self.grad_scaler.unscale_(optimizer)
def step(self, optimizer: Optimizer, *args, **kwargs):
return self.grad_scaler.step(optimizer, *args, **kwargs)
def update(self, new_scale: Union[float, Tensor, None] = None):
return self.grad_scaler.update(new_scale=new_scale)
def clip_grad_norm_(self, params: TensorOrIterableTensors):
torch.nn.utils.clip_grad_norm_(params, self.max_norm)
def state_dict(self) -> dict:
return self.grad_scaler.state_dict()
def load_state_dict(self, state_dict: dict):
return self.grad_scaler.load_state_dict(state_dict)
def __call__(
self,
loss: Tensor,
optimizer: torch.optim.Optimizer,
parameters: Optional[TensorOrIterableTensors] = None,
zero_grad_set_to_none: bool = False,
):
self.scale(loss).backward()
if self.max_norm is not None:
assert parameters is not None
self.unscale_(optimizer)
self.clip_grad_norm_(parameters)
self.grad_scaler.step(optimizer)
self.grad_scaler.update()
optimizer.zero_grad(set_to_none=zero_grad_set_to_none)
def backward(
self,
loss: Tensor,
optimizer: torch.optim.Optimizer,
parameters: Optional[TensorOrIterableTensors] = None,
):
return self(loss, optimizer, parameters=parameters)
def before_train(self):
if self._grad_scaler_config is None:
self._scaler = GradScaler()
else:
self._scaler = GradScaler(**self._grad_scaler_config)
def train(
epoch: int,
data: DistributedDataObject,
device: torch.device,
rank: int,
model: nn.Module,
loss_fn: LossFunction,
optimizer: optim.Optimizer,
args: dict,
scaler: GradScaler = None,
):
model.train()
# Horovod: set epoch to sampler for shuffling
data.sampler.set_epoch(epoch)
running_loss = torch.tensor(0.0)
training_acc = torch.tensor(0.0)
if torch.cuda.is_available():
running_loss = running_loss.to(device)
training_acc = training_acc.to(device)
for batch_idx, (batch, target) in enumerate(data.loader):
if torch.cuda.is_available():
batch, target = batch.to(device), target.to(device)
optimizer.zero_grad()
output = model(batch)
loss = loss_fn(output, target)
if scaler is not None:
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
else:
loss.backward()
optimizer.step()
pred = output.data.max(1, keepdim=True)[1]
acc = pred.eq(target.data.view_as(pred)).cpu().float().sum()
training_acc += acc
running_loss += loss.item()
if batch_idx % args.log_interval == 0:
metrics_ = {
'epoch': epoch,
'batch_loss': loss.item() / args.batch_size,
'running_loss': running_loss / len(data.sampler),
'batch_acc': acc.item() / args.batch_size,
'training_acc': training_acc / len(data.sampler),
}
jdx = batch_idx * len(batch)
frac = 100. * batch_idx / len(data.loader)
pre = [
f'[{rank}]',
f'[{jdx:>5}/{len(data.sampler):<5} ({frac:>03.1f}%)]'
]
io.print_metrics(metrics_, pre=pre, logger=logger)
running_loss = running_loss / len(data.sampler)
training_acc = training_acc / len(data.sampler)
loss_avg = metric_average(running_loss)
training_acc = metric_average(training_acc)
if rank == 0:
logger.log(f'training set; avg loss: {loss_avg:.4g}, '
f'accuracy: {training_acc * 100:.2f}%')
class Trainer:
def __init__(self, config, device='cuda:0'):
self.config = config
self.device = torch.device(device)
self.model_dir = self.config['model_dir']
self.dev_gold_path = os.path.join(self.model_dir, 'dev-gold.txt')
self.dev_pred_path = os.path.join(self.model_dir, 'dev-pred.txt')
self.best_smatch = -1 # May get reloaded if using a checkpoint
self.start_epoch = 1 # May get reloaded if using a checkpoint
os.makedirs(self.model_dir, exist_ok=True)
def train(self, checkpoint=None):
self.load_model(
checkpoint
) # sets self.model, tokenizer, optimizer, .., best_smatch, start_epoch
self.load_training_data() # sets self.train_loader
self.load_eval_data() # sets self.inference, graphs_gold
# Loop through max epochs
assert self.start_epoch < self.config[
'max_epochs'] # May fail if reloading a checkpoint
for epoch in range(self.start_epoch, self.config['max_epochs'] + 1):
# Setup batch
print('Training epoch %d' % epoch)
trn_amr_loss = RunningAverage()
self.optimizer.zero_grad()
pbar = tqdm(total=len(self.train_loader.dataset), ncols=100)
self.set_pbar_desc_train(pbar, None)
self.model.train()
# Loop through all the data
for bnum, batch in enumerate(self.train_loader):
x, y, extra = batch
with autocast(enabled=self.config['fp16']):
rdict = self.model(**x, **y)
loss = rdict['loss']
self.scaler.scale(
(loss / self.config['accum_steps'])).backward()
trn_amr_loss.update(loss.item())
# Perform an update every accum_steps
if (bnum + 1) % self.config['accum_steps'] == 0:
self.step_otimizer()
# Update progress
pbar.update(x['input_ids'].shape[0])
self.set_pbar_desc_train(pbar, trn_amr_loss.value)
pbar.close()
# Perform an update with the last batch if it wasn't already done in the loop
if (bnum + 1) % self.config['accum_steps'] != 0:
self.step_otimizer()
# Run evaluate, compute smatch and save the model if it's the new best
try:
smatch = self.evaluate()
if smatch > self.best_smatch:
self.best_smatch = smatch
self.save_and_remove_checkpoints(epoch, smatch)
except:
print('!! Evaluation / save failed !!')
logger.exception('Evaluation or model save failed')
print()
# Run Inference and evaluate the model
def evaluate(self):
self.model.eval()
sents = [g.metadata['snt'] for g in self.graphs_gold]
graphs_gen = self.inference.parse_sents(sents,
return_penman=True,
disable_progress=False,
pbar_desc='%-14s' %
'Evaluating:')
assert len(graphs_gen) == len(self.graphs_gold)
# Detect bad graphs. In Penman 1.2.0, metadata does not impact penam.Graph.__eq__()
num_bad = sum(g == Inference.invalid_graph for g in graphs_gen)
print('Out of %d graphs, %d did not generate properly.' %
(len(graphs_gen), num_bad))
# Save the final graphs
print('Generated graphs written to', self.dev_pred_path)
penman.dump(graphs_gen, self.dev_pred_path, indent=6, model=amr_model)
# Run smatch
try:
gold_entries = get_entries(self.dev_gold_path)
test_entries = get_entries(self.dev_pred_path)
precision, recall, f_score = compute_smatch(
test_entries, gold_entries)
print('SMATCH -> P: %.3f, R: %.3f, F: %.3f' %
(precision, recall, f_score))
except:
logger.exception('Failed to compute smatch score.')
precision, recall, f_score = 0, 0, 0
return f_score
# Save the checkpoints if this is the best score
def save_and_remove_checkpoints(self, epoch, smatch):
prev_checkpoints = [
fn for fn in os.listdir(self.model_dir) if fn.endswith('.pt')
]
model_fn = 'checkpoint_epoch_%02d_smatch_%04d.pt' % (epoch,
smatch * 10000)
model_fpath = os.path.join(self.model_dir, model_fn)
# Create the dictionary with the optional optimizer and save it
print('Saving new, best model to', model_fpath)
save_dict = {'model': self.model.state_dict()}
if self.config.get('save_optimizer'):
save_dict['optimizer'] = self.optimizer.state_dict()
save_dict['scheduler'] = self.scheduler.state_dict()
torch.save(save_dict, model_fpath)
# Save the config file
self.config['smatch_dev'] = smatch
self.config['last_epoch'] = epoch
with open(os.path.join(self.model_dir, 'config.json'), 'w') as f:
json.dump(self.config, f, indent=4)
# Remove previous checkpoints
for chkpt_fn in prev_checkpoints:
os.remove(os.path.join(self.model_dir, chkpt_fn))
# Load and setup the model, tokenizer, optimizer, etc..
def load_model(self, checkpoint=None):
print('Loading model from', self.config['model'])
self.model, self.tokenizer = instantiate_model_and_tokenizer(
self.config['model'],
additional_tokens_smart_init=self.config['smart_init'],
dropout=self.config['dropout'],
attention_dropout=self.config['attention_dropout'],
penman_linearization=self.config['penman_linearization'],
collapse_name_ops=self.config['collapse_name_ops'],
use_pointer_tokens=self.config['use_pointer_tokens'],
raw_graph=self.config['raw_graph'])
self.model.to(self.device)
# Load optimization components
self.optimizer = AdamW(self.model.parameters(),
lr=self.config['learning_rate'],
weight_decay=self.config['weight_decay'])
self.scheduler = transformers.get_constant_schedule_with_warmup(
self.optimizer, num_warmup_steps=self.config['warmup_steps'])
self.scaler = GradScaler(enabled=self.config['fp16'])
# Reload checkpoint model weights and optimizer params if loading from a checkpoint
if checkpoint is not None:
print('Checkpoint %s restored' % checkpoint)
load_state_dict_from_checkpoint(checkpoint, self.model,
self.optimizer, self.scheduler)
# Try to load the smatch score and last_epoch from the config in the model directory.
try:
with open(os.path.join(self.model_dir, 'config.json')) as f:
model_config = json.load(f)
self.best_smatch = model_config['smatch_dev']
self.start_epoch = model_config['last_epoch'] + 1
except:
logger.exception(
'Unable to load config file in model directory')
# Setup the training data loader
def load_training_data(self):
print('Loading train data from', self.config['train'])
self.train_loader = get_dataloader(
self.tokenizer,
glob_pattern=self.config['train'],
evaluation=False,
batch_size=self.config['batch_size'],
use_recategorization=self.config['use_recategorization'],
remove_longer_than=self.config['remove_longer_than'],
remove_wiki=self.config['remove_wiki'],
dereify=self.config['dereify'],
device=self.device)
# Setup the inference object and create the gold data test file
def load_eval_data(self):
print('Loading eval data from ', self.config['dev'])
self.inference = Inference(model=self.model,
tokenizer=self.tokenizer,
device=self.device,
num_beams=self.config['eval_beam_size'],
batch_size=self.config['eval_batch_sents'],
config=self.config)
self.graphs_gold = read_raw_amr_data(
self.config['dev'],
use_recategorization=self.config['use_recategorization'],
dereify=self.config['dereify'],
remove_wiki=self.config['remove_wiki'])
penman.dump(self.graphs_gold,
self.dev_gold_path,
indent=6,
model=amr_model)
# Function to update the model's parameters for accumulated loss
def step_otimizer(self):
self.scaler.unscale_(self.optimizer)
torch.nn.utils.clip_grad_norm_(self.model.parameters(),
self.config['grad_norm'])
self.scaler.step(self.optimizer)
self.scaler.update()
self.optimizer.zero_grad()
self.scheduler.step()
# Update tqdm progress bar description with loss values
@staticmethod
def set_pbar_desc_train(pbar, av_loss):
desc = 'Loss: '
if av_loss is None:
desc += ' ' * 8
else:
desc += '%8.3f' % av_loss
pbar.set_description(desc)
class MyMixedDefaultTrainer(TrainerBase):
"""w/o AMP mode"""
def __init__(self, cfg):
"""
Args:
cfg (CfgNode):
"""
super().__init__()
logger = logging.getLogger("fastreid")
if not logger.isEnabledFor(
logging.INFO): # setup_logger is not called for fastreid
setup_logger()
# Assume these objects must be constructed in this order.
data_loader = self.build_train_loader(cfg)
cfg = self.auto_scale_hyperparams(cfg, data_loader.dataset.num_classes)
self.model = self.build_model(cfg)
self.optimizer, self.param_wrapper = self.build_optimizer(
cfg, self.model)
# For training, wrap with DDP. But don't need this for inference.
if comm.get_world_size() > 1:
# ref to https://github.com/pytorch/pytorch/issues/22049 to set `find_unused_parameters=True`
# for part of the parameters is not updated.
self.model = DistributedDataParallel(
self.model,
device_ids=[comm.get_local_rank()],
broadcast_buffers=False,
)
self._data_loader_iter = iter(data_loader)
self.iters_per_epoch = len(
data_loader.dataset) // cfg.SOLVER.IMS_PER_BATCH
self.scheduler = self.build_lr_scheduler(cfg, self.optimizer,
self.iters_per_epoch)
# Assume no other objects need to be checkpointed.
# We can later make it checkpoint the stateful hooks
self.checkpointer = Checkpointer(
# Assume you want to save checkpoints together with logs/statistics
self.model,
cfg.OUTPUT_DIR,
save_to_disk=comm.is_main_process(),
optimizer=self.optimizer,
**self.scheduler,
)
self.start_epoch = 0
self.max_epoch = cfg.SOLVER.MAX_EPOCH
self.max_iter = self.max_epoch * self.iters_per_epoch
self.warmup_iters = cfg.SOLVER.WARMUP_ITERS
self.delay_epochs = cfg.SOLVER.DELAY_EPOCHS
self.cfg = cfg
self.register_hooks(self.build_hooks())
if cfg.SOLVER.AMP.ENABLED:
unsupported = f"[{self.__class__.__name__}] does not support single-process multi-device training!"
if isinstance(self.model, DistributedDataParallel):
assert not (self.model.device_ids
and len(self.model.device_ids) > 1), unsupported
from torch.cuda.amp.grad_scaler import GradScaler
self.grad_scaler = GradScaler()
else:
self.grad_scaler = None
def resume_or_load(self, resume=True):
"""
If `resume==True` and `cfg.OUTPUT_DIR` contains the last checkpoint (defined by
a `last_checkpoint` file), resume from the file. Resuming means loading all
available states (eg. optimizer and scheduler) and update iteration counter
from the checkpoint. ``cfg.MODEL.WEIGHTS`` will not be used.
Otherwise, this is considered as an independent training. The method will load model
weights from the file `cfg.MODEL.WEIGHTS` (but will not load other states) and start
from iteration 0.
Args:
resume (bool): whether to do resume or not
"""
# The checkpoint stores the training iteration that just finished, thus we start
# at the next iteration (or iter zero if there's no checkpoint).
checkpoint = self.checkpointer.resume_or_load(self.cfg.MODEL.WEIGHTS,
resume=resume)
if resume and self.checkpointer.has_checkpoint():
self.start_epoch = checkpoint.get("epoch", -1) + 1
# The checkpoint stores the training iteration that just finished, thus we start
# at the next iteration (or iter zero if there's no checkpoint).
def build_hooks(self):
"""
Build a list of default hooks, including timing, evaluation,
checkpointing, lr scheduling, precise BN, writing events.
Returns:
list[HookBase]:
"""
logger = logging.getLogger(__name__)
cfg = self.cfg.clone()
cfg.defrost()
cfg.DATALOADER.NUM_WORKERS = 0 # save some memory and time for PreciseBN
cfg.DATASETS.NAMES = tuple([cfg.TEST.PRECISE_BN.DATASET
]) # set dataset name for PreciseBN
ret = [
hooks.IterationTimer(),
hooks.LRScheduler(self.optimizer, self.scheduler),
]
if cfg.TEST.PRECISE_BN.ENABLED and hooks.get_bn_modules(self.model):
logger.info("Prepare precise BN dataset")
ret.append(
hooks.PreciseBN(
# Run at the same freq as (but before) evaluation.
self.model,
# Build a new data loader to not affect training
self.build_train_loader(cfg),
cfg.TEST.PRECISE_BN.NUM_ITER,
))
if len(cfg.MODEL.FREEZE_LAYERS) > 0 and cfg.SOLVER.FREEZE_ITERS > 0:
ret.append(
hooks.LayerFreeze(
self.model,
cfg.MODEL.FREEZE_LAYERS,
cfg.SOLVER.FREEZE_ITERS,
))
# Do PreciseBN before checkpointer, because it updates the model and need to
# be saved by checkpointer.
# This is not always the best: if checkpointing has a different frequency,
# some checkpoints may have more precise statistics than others.
def test_and_save_results():
self._last_eval_results = self.test(self.cfg, self.model)
return self._last_eval_results
# Do evaluation before checkpointer, because then if it fails,
# we can use the saved checkpoint to debug.
ret.append(hooks.EvalHook(cfg.TEST.EVAL_PERIOD, test_and_save_results))
if comm.is_main_process():
ret.append(
hooks.PeriodicCheckpointer(self.checkpointer,
cfg.SOLVER.CHECKPOINT_PERIOD))
# run writers in the end, so that evaluation metrics are written
ret.append(hooks.PeriodicWriter(self.build_writers(), 200))
return ret
def build_writers(self):
"""
Build a list of writers to be used. By default it contains
writers that write metrics to the screen,
a json file, and a tensorboard event file respectively.
If you'd like a different list of writers, you can overwrite it in
your trainer.
Returns:
list[EventWriter]: a list of :class:`EventWriter` objects.
It is now implemented by:
.. code-block:: python
return [
CommonMetricPrinter(self.max_iter),
JSONWriter(os.path.join(self.cfg.OUTPUT_DIR, "metrics.json")),
TensorboardXWriter(self.cfg.OUTPUT_DIR),
]
"""
# Assume the default print/log frequency.
return [
# It may not always print what you want to see, since it prints "common" metrics only.
CommonMetricPrinter(self.max_iter),
JSONWriter(os.path.join(self.cfg.OUTPUT_DIR, "metrics.json")),
TensorboardXWriter(self.cfg.OUTPUT_DIR),
]
def train(self):
"""
Run training.
Returns:
OrderedDict of results, if evaluation is enabled. Otherwise None.
"""
super().train(self.start_epoch, self.max_epoch, self.iters_per_epoch)
if comm.is_main_process():
assert hasattr(self, "_last_eval_results"
), "No evaluation results obtained during training!"
return self._last_eval_results
def run_step(self):
assert self.model.training, f"[{self.__class__.__name__}] model was changed to eval mode!"
if self.cfg.SOLVER.AMP.ENABLED:
assert torch.cuda.is_available(
), f"[{self.__class__.__name__}] CUDA is required for AMP training!"
from torch.cuda.amp.autocast_mode import autocast
start = time.perf_counter()
data = next(self._data_loader_iter)
data_time = time.perf_counter() - start
if self.cfg.SOLVER.AMP.ENABLED:
with autocast():
loss_dict = self.model(data)
losses = sum(loss_dict.values())
self.optimizer.zero_grad()
self.grad_scaler.scale(losses).backward()
self._write_metrics(loss_dict, data_time)
self.grad_scaler.step(self.optimizer)
self.grad_scaler.update()
else:
loss_dict = self.model(data)
losses = sum(loss_dict.values())
self.optimizer.zero_grad()
losses.backward()
self._write_metrics(loss_dict, data_time)
self.optimizer.step()
if isinstance(self.param_wrapper, ContiguousParams):
self.param_wrapper.assert_buffer_is_valid()
def _write_metrics(self, loss_dict: Dict[str, torch.Tensor],
data_time: float):
"""
Args:
loss_dict (dict): dict of scalar losses
data_time (float): time taken by the dataloader iteration
"""
device = next(iter(loss_dict.values())).device
# Use a new stream so these ops don't wait for DDP or backward
with torch.cuda.stream(torch.cuda.Stream() if device.type ==
"cuda" else None):
metrics_dict = {
k: v.detach().cpu().item()
for k, v in loss_dict.items()
}
metrics_dict["data_time"] = data_time
# Gather metrics among all workers for logging
# This assumes we do DDP-style training, which is currently the only
# supported method in detectron2.
all_metrics_dict = comm.gather(metrics_dict)
if comm.is_main_process():
storage = get_event_storage()
# data_time among workers can have high variance. The actual latency
# caused by data_time is the maximum among workers.
data_time = np.max([x.pop("data_time") for x in all_metrics_dict])
storage.put_scalar("data_time", data_time)
# average the rest metrics
metrics_dict = {
k: np.mean([x[k] for x in all_metrics_dict])
for k in all_metrics_dict[0].keys()
}
total_losses_reduced = sum(metrics_dict.values())
if not np.isfinite(total_losses_reduced):
raise FloatingPointError(
f"Loss became infinite or NaN at iteration={self.iter}!\n"
f"loss_dict = {metrics_dict}")
storage.put_scalar("total_loss", total_losses_reduced)
if len(metrics_dict) > 1:
storage.put_scalars(**metrics_dict)
@classmethod
def build_model(cls, cfg, show_model=True):
"""
Returns:
torch.nn.Module:
It now calls :func:`fastreid.modeling.build_model`.
Overwrite it if you'd like a different model.
"""
model = build_model(cfg)
if show_model:
logger = logging.getLogger('fastreid')
logger.info("Model:\n{}".format(model))
return model
@classmethod
def build_optimizer(cls, cfg, model):
"""
Returns:
torch.optim.Optimizer:
It now calls :func:`fastreid.solver.build_optimizer`.
Overwrite it if you'd like a different optimizer.
"""
return build_optimizer(cfg, model)
@classmethod
def build_lr_scheduler(cls, cfg, optimizer, iters_per_epoch):
"""
It now calls :func:`fastreid.solver.build_lr_scheduler`.
Overwrite it if you'd like a different scheduler.
"""
return build_lr_scheduler(cfg, optimizer, iters_per_epoch)
@classmethod
def build_train_loader(cls, cfg):
"""
Returns:
iterable
It now calls :func:`fastreid.data.build_reid_train_loader`.
Overwrite it if you'd like a different data loader.
"""
logger = logging.getLogger(__name__)
logger.info("Prepare training set")
return build_reid_train_loader(cfg, combineall=cfg.DATASETS.COMBINEALL)
@classmethod
def build_test_loader(cls, cfg, dataset_name):
"""
Returns:
iterable
It now calls :func:`fastreid.data.build_reid_test_loader`.
Overwrite it if you'd like a different data loader.
"""
return build_reid_test_loader(cfg, dataset_name=dataset_name)
@classmethod
def build_evaluator(cls, cfg, dataset_name, output_dir=None):
data_loader, num_query = cls.build_test_loader(cfg, dataset_name)
return data_loader, ReidEvaluator(cfg, num_query, output_dir)
@classmethod
def test(cls, cfg, model):
"""
Args:
cfg (CfgNode):
model (nn.Module):
Returns:
dict: a dict of result metrics
"""
logger = logging.getLogger(__name__)
results = OrderedDict()
for idx, dataset_name in enumerate(cfg.DATASETS.TESTS):
logger.info("Prepare testing set")
try:
data_loader, evaluator = cls.build_evaluator(cfg, dataset_name)
except NotImplementedError:
logger.warn(
"No evaluator found. implement its `build_evaluator` method."
)
results[dataset_name] = {}
continue
results_i = inference_on_dataset(model,
data_loader,
evaluator,
flip_test=cfg.TEST.FLIP.ENABLED)
results[dataset_name] = results_i
if comm.is_main_process():
assert isinstance(
results, dict
), "Evaluator must return a dict on the main process. Got {} instead.".format(
results)
logger.info("Evaluation results for {} in csv format:".format(
dataset_name))
results_i['dataset'] = dataset_name
print_csv_format(results_i)
if len(results) == 1:
results = list(results.values())[0]
return results
@staticmethod
def auto_scale_hyperparams(cfg, num_classes):
r"""
This is used for auto-computation actual training iterations,
because some hyper-param, such as MAX_ITER, means training epochs rather than iters,
so we need to convert specific hyper-param to training iterations.
"""
cfg = cfg.clone()
frozen = cfg.is_frozen()
cfg.defrost()
# If you don't hard-code the number of classes, it will compute the number automatically
if cfg.MODEL.HEADS.NUM_CLASSES == 0:
output_dir = cfg.OUTPUT_DIR
cfg.MODEL.HEADS.NUM_CLASSES = num_classes
logger = logging.getLogger(__name__)
logger.info(
f"Auto-scaling the num_classes={cfg.MODEL.HEADS.NUM_CLASSES}")
# Update the saved config file to make the number of classes valid
if comm.is_main_process() and output_dir:
# Note: some of our scripts may expect the existence of
# config.yaml in output directory
path = os.path.join(output_dir, "config.yaml")
with PathManager.open(path, "w") as f:
f.write(cfg.dump())
if frozen: cfg.freeze()
return cfg
class Trainer:
def __init__(self, cfg):
self.cfg = cfg
self.paths = cfg['paths']
self.net_params = cfg['net']
self.train_params = cfg['train']
self.trans_params = cfg['train']['transforms']
self.checkpoints = self.paths['checkpoints']
Path(self.checkpoints).mkdir(parents=True, exist_ok=True)
shutil.copyfile('config.yaml', f'{self.checkpoints}/config.yaml')
self.update_interval = self.paths['update_interval']
# amp training
self.use_amp = self.train_params['mixed_precision']
self.scaler = GradScaler() if self.use_amp else None
# data setup
dataset_name = self.train_params['dataset']
self.use_multi = dataset_name == 'multi'
print(f'Using dataset: {dataset_name}')
self.train_dataset = get_pedestrian_dataset(
dataset_name,
self.paths,
augment=get_train_transforms(self.trans_params),
mode='train',
multi_datasets=self.train_params['multi_datasets']
if self.use_multi else None)
print(f'Train dataset: {len(self.train_dataset)} samples')
self.val_dataset = get_pedestrian_dataset(
dataset_name,
self.paths,
augment=get_val_transforms(self.trans_params),
mode='val',
multi_datasets=self.train_params['multi_datasets']
if self.use_multi else None)
print(f'Val dataset: {len(self.val_dataset)} samples')
tests_data = self.train_params['test_datasets']
self.test_datasets = [
get_pedestrian_dataset(d_name,
self.paths,
augment=get_test_transforms(
self.trans_params),
mode='test') for d_name in tests_data
]
self.criterion = AnchorFreeLoss(self.train_params)
self.writer = Writer(self.paths['log_dir'])
print('Tensorboard logs are saved to: {}'.format(
self.paths['log_dir']))
self.sched_type = self.train_params['scheduler']
self.scheduler = None
self.optimizer = None
def save_checkpoints(self, epoch, net):
path = osp.join(self.checkpoints, f'Epoch_{epoch}.pth')
torch.save(
{
'epoch': epoch,
'net_state_dict': net.state_dict(),
'optimizer_state_dict': self.optimizer.state_dict()
}, path)
def train(self):
torch.backends.cudnn.benchmark = True
torch.backends.cudnn.deterministic = True
batch_size = self.train_params['batch_size']
self.batch_size = batch_size
num_workers = self.train_params['num_workers']
pin_memory = self.train_params['pin_memory']
print('Batch-size = {}'.format(batch_size))
train_loader = DataLoader(self.train_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=num_workers,
pin_memory=pin_memory,
drop_last=True)
val_loader = DataLoader(self.val_dataset,
batch_size=batch_size,
shuffle=False,
num_workers=num_workers,
pin_memory=pin_memory,
drop_last=False)
# net setup
print('Preparing net: ')
net = get_fpn_net(self.net_params)
# train setup
lr = self.train_params['lr']
epochs = self.train_params['epochs']
weight_decay = self.train_params['weight_decay']
self.optimizer = optim.Adam(net.parameters(),
lr=lr,
weight_decay=weight_decay,
eps=1e-4)
if self.net_params['pretrained']:
checkpoint = torch.load(self.net_params['pretrained_model'],
map_location="cuda")
net.load_state_dict(checkpoint['net_state_dict'])
self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
for p in self.optimizer.param_groups:
p['lr'] = lr
for state in self.optimizer.state.values():
for k, v in state.items():
if torch.is_tensor(v):
state[k] = v.cuda()
print('CHECKPOINT LOADED')
net.cuda()
first_epoch = 0
# scheduler
if self.sched_type == 'ocp':
last_epoch = -1 if first_epoch == 0 else first_epoch * len(
train_loader)
self.scheduler = OneCycleLR(
self.optimizer,
max_lr=lr,
epochs=epochs,
last_epoch=last_epoch,
steps_per_epoch=len(train_loader),
pct_start=self.train_params['ocp_params']['max_lr_pct'])
elif self.sched_type == 'multi_step':
last_epoch = -1 if first_epoch == 0 else first_epoch
self.scheduler = MultiStepLR(
self.optimizer,
milestones=self.train_params['multi_params']['milestones'],
gamma=self.train_params['multi_params']['gamma'],
last_epoch=last_epoch)
#start training
net.train()
val_rate = self.train_params['val_rate']
test_rate = self.train_params['test_rate']
for epoch in range(first_epoch, epochs):
self.train_epoch(net, train_loader, epoch)
if self.sched_type != 'ocp':
self.writer.log_lr(epoch, self.scheduler.get_last_lr()[0])
self.scheduler.step()
if (epoch + 1) % val_rate == 0 or epoch == epochs - 1:
self.eval(net, val_loader, epoch * len(train_loader))
if (epoch + 1) % (val_rate *
test_rate) == 0 or epoch == epochs - 1:
self.test_ap(net, epoch)
self.save_checkpoints(epoch, net)
def train_epoch(self, net, loader, epoch):
net.train()
loss_metric = LossMetric(self.cfg)
probs = ProbsAverageMeter()
for mini_batch_i, read_mini_batch in tqdm(enumerate(loader),
desc=f'Epoch {epoch}:',
ascii=True,
total=len(loader)):
data, labels = read_mini_batch
data = data.cuda()
labels = [label.cuda() for label in labels]
with amp.autocast():
out = net(data)
loss_dict, hm_probs = self.criterion(out, labels)
loss = loss_metric.calculate_loss(loss_dict)
self.optimizer.zero_grad()
self.scaler.scale(loss).backward()
self.scaler.step(self.optimizer)
self.scaler.update()
probs.update(hm_probs)
if self.sched_type == 'ocp':
self.scheduler.step()
loss_metric.add_sample(loss_dict)
if mini_batch_i % self.update_interval == 0:
if self.sched_type == 'ocp':
# TODO write average lr
self.writer.log_lr(epoch * len(loader) + mini_batch_i,
self.scheduler.get_last_lr()[0])
self.writer.log_training(epoch * len(loader) + mini_batch_i,
loss_metric)
self.writer.log_probs(epoch, probs.get_average())
def eval(self, net, loader, step):
net.eval()
loss_metric = LossMetric(self.cfg)
with torch.no_grad():
for _, read_mini_batch in tqdm(enumerate(loader),
desc=f'Val:',
ascii=True,
total=len(loader)):
data, labels = read_mini_batch
data = data.cuda()
labels = [label.cuda() for label in labels]
with amp.autocast():
out = net(data)
loss_dict, _ = self.criterion(out, labels)
loss_metric.add_sample(loss_dict)
self.writer.log_eval(step, loss_metric)
def test_ap(self, net, epoch):
for dataset in self.test_datasets:
ap, _ = test(net, dataset, batch_size=self.batch_size)
self.writer.log_ap(epoch, ap, dataset.name())
def __init__(self, cfg):
"""
Args:
cfg (CfgNode):
"""
super().__init__()
logger = logging.getLogger("fastreid")
if not logger.isEnabledFor(
logging.INFO): # setup_logger is not called for fastreid
setup_logger()
# Assume these objects must be constructed in this order.
data_loader = self.build_train_loader(cfg)
cfg = self.auto_scale_hyperparams(cfg, data_loader.dataset.num_classes)
self.model = self.build_model(cfg)
self.optimizer, self.param_wrapper = self.build_optimizer(
cfg, self.model)
# For training, wrap with DDP. But don't need this for inference.
if comm.get_world_size() > 1:
# ref to https://github.com/pytorch/pytorch/issues/22049 to set `find_unused_parameters=True`
# for part of the parameters is not updated.
self.model = DistributedDataParallel(
self.model,
device_ids=[comm.get_local_rank()],
broadcast_buffers=False,
)
self._data_loader_iter = iter(data_loader)
self.iters_per_epoch = len(
data_loader.dataset) // cfg.SOLVER.IMS_PER_BATCH
self.scheduler = self.build_lr_scheduler(cfg, self.optimizer,
self.iters_per_epoch)
# Assume no other objects need to be checkpointed.
# We can later make it checkpoint the stateful hooks
self.checkpointer = Checkpointer(
# Assume you want to save checkpoints together with logs/statistics
self.model,
cfg.OUTPUT_DIR,
save_to_disk=comm.is_main_process(),
optimizer=self.optimizer,
**self.scheduler,
)
self.start_epoch = 0
self.max_epoch = cfg.SOLVER.MAX_EPOCH
self.max_iter = self.max_epoch * self.iters_per_epoch
self.warmup_iters = cfg.SOLVER.WARMUP_ITERS
self.delay_epochs = cfg.SOLVER.DELAY_EPOCHS
self.cfg = cfg
self.register_hooks(self.build_hooks())
if cfg.SOLVER.AMP.ENABLED:
unsupported = f"[{self.__class__.__name__}] does not support single-process multi-device training!"
if isinstance(self.model, DistributedDataParallel):
assert not (self.model.device_ids
and len(self.model.device_ids) > 1), unsupported
from torch.cuda.amp.grad_scaler import GradScaler
self.grad_scaler = GradScaler()
else:
self.grad_scaler = None
class UDA_Baseline_Trainer(TrainerBase):
"""
load a model pretrained on the source domain,
neglect outliers during training on the target domain
"""
def __init__(self, cfg):
super().__init__()
logger = logging.getLogger("fastreid")
if not logger.isEnabledFor(
logging.INFO): # if setup_logger is not called for fastreid
setup_logger()
logger.info("==> Load target-domain dataset")
self.tgt = tgt = self.load_dataset(cfg.DATASETS.TGT)
self.tgt_nums = len(tgt.train)
cfg = self.auto_scale_hyperparams(cfg, self.tgt_nums)
# Create model
self.model = self.build_model(cfg,
load_model=cfg.MODEL.PRETRAIN,
show_model=True,
use_dsbn=False)
# Optimizer
self.optimizer, self.param_wrapper = self.build_optimizer(
cfg, self.model)
# For training, wrap with DDP. But don't need this for inference.
if comm.get_world_size() > 1:
# ref to https://github.com/pytorch/pytorch/issues/22049 to set `find_unused_parameters=True`
# for part of the parameters is not updated.
self.model = DistributedDataParallel(
self.model,
device_ids=[comm.get_local_rank()],
broadcast_buffers=False,
find_unused_parameters=True)
# Learning rate scheduler
self.iters_per_epoch = cfg.SOLVER.ITERS
self.scheduler = self.build_lr_scheduler(cfg, self.optimizer,
self.iters_per_epoch)
# Assume no other objects need to be checkpointed.
# We can later make it checkpoint the stateful hooks
self.checkpointer = Checkpointer(
# Assume you want to save checkpoints together with logs/statistics
self.model,
cfg.OUTPUT_DIR,
save_to_disk=comm.is_main_process(),
optimizer=self.optimizer,
**self.scheduler,
)
self.start_epoch = 0
self.max_epoch = cfg.SOLVER.MAX_EPOCH
self.max_iter = self.max_epoch * self.iters_per_epoch
self.warmup_iters = cfg.SOLVER.WARMUP_ITERS
self.delay_epochs = cfg.SOLVER.DELAY_EPOCHS
self.cfg = cfg
self.register_hooks(self.build_hooks())
if cfg.SOLVER.AMP.ENABLED:
unsupported = "AMPTrainer does not support single-process multi-device training!"
if isinstance(self.model, DistributedDataParallel):
assert not (self.model.device_ids
and len(self.model.device_ids) > 1), unsupported
from torch.cuda.amp.grad_scaler import GradScaler
self.grad_scaler = GradScaler()
else:
self.grad_scaler = None
def train(self):
"""
Run training.
Returns:
OrderedDict of results, if evaluation is enabled. Otherwise None.
"""
super().train(self.start_epoch, self.max_epoch, self.iters_per_epoch)
if comm.is_main_process():
assert hasattr(self, "_last_eval_results"
), "No evaluation results obtained during training!"
return self._last_eval_results
def before_train(self):
self.model.train()
if self.cfg.SOLVER.AMP.ENABLED:
assert torch.cuda.is_available(
), "CUDA is required for AMP training!"
return super().before_train()
def before_epoch(self):
logger = logging.getLogger('fastreid')
# Calculate distance
logger.info("==> Create pseudo labels for unlabeled target domain")
with inference_context(self.model), torch.no_grad():
tgt_train = self.build_dataset(self.cfg,
self.tgt.train,
is_train=False,
relabel=False,
with_mem_idx=False)
tgt_init_feat_loader = self.build_test_loader(self.cfg, tgt_train)
tgt_fname_feat_dict, _ = extract_features(self.model,
tgt_init_feat_loader)
tgt_features = torch.cat([
tgt_fname_feat_dict[f].unsqueeze(0)
for f, _, _ in sorted(self.tgt.train)
], 0)
tgt_features = F.normalize(tgt_features, dim=1)
rerank_dist = compute_jaccard_distance(tgt_features,
k1=self.cfg.CLUSTER.JACCARD.K1,
k2=self.cfg.CLUSTER.JACCARD.K2)
if self.epoch == 0:
if self.cfg.CLUSTER.DBSCAN.ADAPTIVE_EPS:
logger.info("==> Calculating eps according to rerank_dist...")
tri_mat = np.triu(rerank_dist, 1) # tri_mat.dim=2
tri_mat = tri_mat[np.nonzero(tri_mat)] # tri_mat.dim=1
tri_mat = np.sort(tri_mat, axis=None)
top_num = np.round(self.cfg.SOLVER.RHO *
tri_mat.size).astype(int)
self.eps = tri_mat[:top_num].mean()
logger.info(f"==> epoch {self.epoch} eps: {self.eps}")
else:
self.eps = self.cfg.CLUSTER.DBSCAN.EPS
self.cluster = DBSCAN(eps=self.eps,
min_samples=4,
metric="precomputed",
n_jobs=-1)
# select & cluster images as training set of this epochs
logger.info(f"Clustering and labeling...")
pseudo_labels = self.cluster.fit_predict(rerank_dist)
self.num_clusters = num_clusters = len(
set(pseudo_labels)) - (1 if -1 in pseudo_labels else 0)
num_outliers = pseudo_labels[pseudo_labels == -1].shape[0]
# pseudo_labels = self.generate_pseudo_labels(pseudo_labels, num_clusters)
# pseudo_labels = self.assign_outlier(pseudo_labels, tgt_features)
del tgt_features
pseudo_labeled_dataset = []
cluster_centers = collections.defaultdict(list)
for i, ((fname, _, cid),
label) in enumerate(zip(sorted(self.tgt.train),
pseudo_labels)):
if label != -1:
pseudo_labeled_dataset.append((fname, label, cid))
cluster_centers[label].append(tgt_fname_feat_dict[fname])
del tgt_fname_feat_dict, rerank_dist
cluster_centers = [
torch.stack(cluster_centers[idx]).mean(0)
for idx in sorted(cluster_centers.keys())
]
cluster_centers = torch.stack(cluster_centers)
if isinstance(self.model, DistributedDataParallel):
self.model.module.heads.weight.data[:num_clusters].copy_(
F.normalize(cluster_centers, dim=1).float().cuda())
else:
self.model.heads.weight.data[:num_clusters].copy_(
F.normalize(cluster_centers, dim=1).float().cuda())
# statistics of clusters and un-clustered instances
# index2label = collections.defaultdict(int)
# for label in pseudo_labels:
# index2label[label.item()] += 1
# print(f'cluster_label', min(cluster_label), max(cluster_label), len(cluster_label))
# print(f'outlier label', min(outlier_label), max(outlier_label), len(outlier_label))
# index2label = np.fromiter(index2label.values(), dtype=float)
logger.info(
"==> Statistics for epoch {}: {} clusters, {} un-clustered instances"
.format(self.epoch, num_clusters, num_outliers))
pseudo_tgt_train = self.build_dataset(
self.cfg,
pseudo_labeled_dataset,
is_train=True,
relabel=True, # relabel?
# relabel=False,
with_mem_idx=False)
self.pseudo_tgt_train_loader = self.build_train_loader(
self.cfg,
train_set=pseudo_tgt_train,
sampler=RandomMultipleGallerySampler(
pseudo_tgt_train.img_items,
self.cfg.SOLVER.IMS_PER_BATCH // comm.get_world_size(),
self.cfg.DATALOADER.NUM_INSTANCE),
with_mem_idx=False)
return super().before_epoch()
# assign outlier to its nearest neighbor
def assign_outlier(self, pseudo_labels, tgt_features):
outlier_mask = pseudo_labels == -1
cluster_mask = pseudo_labels != -1
outlier_feat = tgt_features[outlier_mask]
cluster_feat = tgt_features[cluster_mask]
dist = torch.cdist(outlier_feat, cluster_feat) # 计算特征间L2距离
min_dist_idx = dist.argmin(-1).cpu()
pseudo_labels[outlier_mask] = pseudo_labels[cluster_mask][min_dist_idx]
return pseudo_labels
def run_step(self):
assert self.model.training, f"[{self.__class__.__name__}] model was changed to eval mode!"
if self.cfg.SOLVER.AMP.ENABLED:
assert torch.cuda.is_available(
), f"[{self.__class__.__name__}] CUDA is required for AMP training!"
from torch.cuda.amp.autocast_mode import autocast
start = time.perf_counter()
# load data
tgt_inputs = self.pseudo_tgt_train_loader.next()
def _parse_data(inputs):
imgs, _, pids, _ = inputs
return imgs.cuda(), pids.cuda()
# process inputs
t_inputs, t_targets = _parse_data(tgt_inputs)
data_time = time.perf_counter() - start
def _forward():
outputs = self.model(t_inputs)
f_out_t = outputs['features']
p_out_t = outputs['pred_class_logits'][:, :self.num_clusters]
loss_dict = {}
loss_ce = cross_entropy_loss(pred_class_outputs=p_out_t,
gt_classes=t_targets,
eps=self.cfg.MODEL.LOSSES.CE.EPSILON,
alpha=self.cfg.MODEL.LOSSES.CE.ALPHA)
loss_dict.update({'loss_ce': loss_ce})
if 'TripletLoss' in self.cfg.MODEL.LOSSES.NAME:
loss_tri = triplet_loss(f_out_t,
t_targets,
margin=0.0,
norm_feat=True,
hard_mining=False)
loss_dict.update({'loss_tri': loss_tri})
return loss_dict
if self.cfg.SOLVER.AMP.ENABLED:
with autocast():
loss_dict = _forward()
losses = sum(loss_dict.values())
self.optimizer.zero_grad()
self.grad_scaler.scale(losses).backward()
self._write_metrics(loss_dict, data_time)
self.grad_scaler.step(self.optimizer)
self.grad_scaler.update()
else:
loss_dict = _forward()
losses = sum(loss_dict.values())
self.optimizer.zero_grad()
losses.backward()
self._write_metrics(loss_dict, data_time)
self.optimizer.step()
if isinstance(self.param_wrapper, ContiguousParams):
self.param_wrapper.assert_buffer_is_valid()
@classmethod
def load_dataset(cls, name):
logger = logging.getLogger(__name__)
logger.info(f"Preparing {name}")
_root = os.getenv("FASTREID_DATASETS", "/root/datasets")
data = DATASET_REGISTRY.get(name)(root=_root)
if comm.is_main_process():
data.show_train()
return data
@classmethod
def build_dataset(cls,
cfg,
img_items,
is_train=False,
relabel=False,
transforms=None,
with_mem_idx=False):
if transforms is None:
transforms = build_transforms(cfg, is_train=is_train)
if with_mem_idx:
sorted_img_items = sorted(img_items)
for i in range(len(sorted_img_items)):
sorted_img_items[i] += (i, )
return InMemoryDataset(sorted_img_items, transforms, relabel)
else:
return CommDataset(img_items, transforms, relabel)
@classmethod
def build_train_loader(cls,
cfg,
train_set=None,
sampler=None,
with_mem_idx=False):
logger = logging.getLogger('fastreid')
logger.info("Prepare training loader")
total_batch_size = cfg.SOLVER.IMS_PER_BATCH
mini_batch_size = total_batch_size // comm.get_world_size()
if sampler is None:
num_instance = cfg.DATALOADER.NUM_INSTANCE
sampler_name = cfg.DATALOADER.SAMPLER_TRAIN
logger.info("Using training sampler {}".format(sampler_name))
if sampler_name == "TrainingSampler":
sampler = samplers.TrainingSampler(len(train_set))
elif sampler_name == "NaiveIdentitySampler":
sampler = samplers.NaiveIdentitySampler(
train_set.img_items, mini_batch_size, num_instance)
elif sampler_name == "BalancedIdentitySampler":
sampler = samplers.BalancedIdentitySampler(
train_set.img_items, mini_batch_size, num_instance)
elif sampler_name == "SetReWeightSampler":
set_weight = cfg.DATALOADER.SET_WEIGHT
sampler = samplers.SetReWeightSampler(train_set.img_items,
mini_batch_size,
num_instance, set_weight)
elif sampler_name == "ImbalancedDatasetSampler":
sampler = samplers.ImbalancedDatasetSampler(
train_set.img_items)
else:
raise ValueError(
"Unknown training sampler: {}".format(sampler_name))
iters = cfg.SOLVER.ITERS
num_workers = cfg.DATALOADER.NUM_WORKERS
batch_sampler = BatchSampler(sampler, mini_batch_size, True)
train_loader = IterLoader(
DataLoader(
Preprocessor(train_set, with_mem_idx),
num_workers=num_workers,
batch_sampler=batch_sampler,
pin_memory=True,
),
length=iters,
)
# train_loader = DataLoaderX(
# comm.get_local_rank(),
# dataset=Preprocessor(train_set, with_mem_idx),
# num_workers=num_workers,
# batch_sampler=batch_sampler,
# collate_fn=fast_batch_collator,
# pin_memory=True,
# )
return train_loader
@classmethod
def build_test_loader(cls, cfg, test_set):
logger = logging.getLogger('fastreid')
logger.info("Prepare testing loader")
# test_loader = DataLoader(
# # Preprocessor(test_set),
# test_set,
# batch_size=cfg.TEST.IMS_PER_BATCH,
# num_workers=cfg.DATALOADER.NUM_WORKERS,
# shuffle=False,
# pin_memory=True,
# )
test_batch_size = cfg.TEST.IMS_PER_BATCH
mini_batch_size = test_batch_size // comm.get_world_size()
num_workers = cfg.DATALOADER.NUM_WORKERS
data_sampler = samplers.InferenceSampler(len(test_set))
batch_sampler = BatchSampler(data_sampler, mini_batch_size, False)
test_loader = DataLoaderX(
comm.get_local_rank(),
dataset=test_set,
batch_sampler=batch_sampler,
num_workers=num_workers, # save some memory
collate_fn=fast_batch_collator,
pin_memory=True,
)
return test_loader
@classmethod
def build_evaluator(cls, cfg, dataset_name, output_dir=None):
data_loader, num_query = build_reid_test_loader(
cfg, dataset_name=dataset_name)
return data_loader, ReidEvaluator(cfg, num_query, output_dir)
@classmethod
def test(cls, cfg, model):
"""
Args:
cfg (CfgNode):
model (nn.Module):
Returns:
dict: a dict of result metrics
"""
logger = logging.getLogger('fastreid')
results = OrderedDict()
dataset_name = cfg.DATASETS.TGT
logger.info("Prepare testing set")
try:
data_loader, evaluator = cls.build_evaluator(cfg, dataset_name)
except NotImplementedError:
logger.warn(
"No evaluator found. implement its `build_evaluator` method.")
results[dataset_name] = {}
results_i = inference_on_dataset(model,
data_loader,
evaluator,
flip_test=cfg.TEST.FLIP.ENABLED)
results[dataset_name] = results_i
if comm.is_main_process():
assert isinstance(
results, dict
), "Evaluator must return a dict on the main process. Got {} instead.".format(
results)
logger.info("Evaluation results for {} in csv format:".format(
dataset_name))
results_i['dataset'] = dataset_name
print_csv_format(results_i)
# if len(results) == 1:
# results = list(results.values())[0]
return results
@classmethod
def build_model(cls,
cfg,
load_model=True,
show_model=True,
use_dsbn=False):
cfg = cfg.clone() # cfg can be modified by model
cfg.defrost()
cfg.MODEL.DEVICE = "cpu"
model = build_model(cfg)
logger = logging.getLogger('fastreid')
if load_model:
pretrain_path = cfg.MODEL.PRETRAIN_PATH
try:
state_dict = torch.load(
pretrain_path, map_location=torch.device("cpu"))['model']
for layer in cfg.MODEL.IGNORE_LAYERS:
if layer in state_dict.keys():
del state_dict[layer]
logger.info(f"Loading pretrained model from {pretrain_path}")
except FileNotFoundError as e:
logger.info(
f"{pretrain_path} is not found! Please check this path.")
raise e
except KeyError as e:
logger.info(
"State dict keys error! Please check the state dict.")
raise e
incompatible = model.load_state_dict(state_dict, strict=False)
if incompatible.missing_keys:
logger.info(
get_missing_parameters_message(incompatible.missing_keys))
if incompatible.unexpected_keys:
logger.info(
get_unexpected_parameters_message(
incompatible.unexpected_keys))
if use_dsbn:
logger.info("==> Convert BN to Domain Specific BN")
convert_dsbn(model)
if show_model:
logger.info("Model:\n{}".format(model))
model.to(torch.device("cuda"))
return model
@staticmethod
def auto_scale_hyperparams(cfg, num_classes):
r"""
This is used for auto-computation actual training iterations,
because some hyper-param, such as MAX_ITER, means training epochs rather than iters,
so we need to convert specific hyper-param to training iterations.
"""
cfg = cfg.clone()
frozen = cfg.is_frozen()
cfg.defrost()
# If you don't hard-code the number of classes, it will compute the number automatically
if cfg.MODEL.HEADS.NUM_CLASSES == 0:
output_dir = cfg.OUTPUT_DIR
cfg.MODEL.HEADS.NUM_CLASSES = num_classes
logger = logging.getLogger('fastreid')
logger.info(
f"Auto-scaling the num_classes={cfg.MODEL.HEADS.NUM_CLASSES}")
# Update the saved config file to make the number of classes valid
if comm.is_main_process() and output_dir:
# Note: some of our scripts may expect the existence of
# config.yaml in output directory
path = os.path.join(output_dir, "config.yaml")
with PathManager.open(path, "w") as f:
f.write(cfg.dump())
if frozen: cfg.freeze()
return cfg
@classmethod
def build_optimizer(cls, cfg, model):
"""
Returns:
torch.optim.Optimizer:
It now calls :func:`fastreid.solver.build_optimizer`.
Overwrite it if you'd like a different optimizer.
"""
return build_optimizer(cfg, model)
@classmethod
def build_lr_scheduler(cls, cfg, optimizer, iters_per_epoch):
"""
It now calls :func:`fastreid.solver.build_lr_scheduler`.
Overwrite it if you'd like a different scheduler.
"""
return build_lr_scheduler(cfg, optimizer, iters_per_epoch)
def build_hooks(self):
"""
Build a list of default hooks, including timing, evaluation,
checkpointing, lr scheduling, precise BN, writing events.
Returns:
list[HookBase]:
"""
logger = logging.getLogger(__name__)
cfg = self.cfg.clone()
cfg.defrost()
cfg.DATALOADER.NUM_WORKERS = 0 # save some memory and time for PreciseBN
cfg.DATASETS.NAMES = tuple([cfg.TEST.PRECISE_BN.DATASET
]) # set dataset name for PreciseBN
ret = [
hooks.IterationTimer(),
hooks.LRScheduler(self.optimizer, self.scheduler),
]
# if cfg.TEST.PRECISE_BN.ENABLED and hooks.get_bn_modules(self.model):
# logger.info("Prepare precise BN dataset")
# ret.append(hooks.PreciseBN(
# # Run at the same freq as (but before) evaluation.
# self.model,
# # Build a new data loader to not affect training
# self.build_train_loader(cfg),
# cfg.TEST.PRECISE_BN.NUM_ITER,
# ))
if len(cfg.MODEL.FREEZE_LAYERS) > 0 and cfg.SOLVER.FREEZE_ITERS > 0:
ret.append(
hooks.LayerFreeze(
self.model,
cfg.MODEL.FREEZE_LAYERS,
cfg.SOLVER.FREEZE_ITERS,
))
# Do PreciseBN before checkpointer, because it updates the model and need to
# be saved by checkpointer.
# This is not always the best: if checkpointing has a different frequency,
# some checkpoints may have more precise statistics than others.
def test_and_save_results():
self._last_eval_results = self.test(self.cfg, self.model)
return self._last_eval_results
# Do evaluation before checkpointer, because then if it fails,
# we can use the saved checkpoint to debug.
ret.append(hooks.EvalHook(cfg.TEST.EVAL_PERIOD, test_and_save_results))
if comm.is_main_process():
ret.append(
hooks.PeriodicCheckpointer(self.checkpointer,
cfg.SOLVER.CHECKPOINT_PERIOD))
# run writers in the end, so that evaluation metrics are written
ret.append(hooks.PeriodicWriter(self.build_writers(), 200))
return ret
def build_writers(self):
"""
Build a list of writers to be used. By default it contains
writers that write metrics to the screen,
a json file, and a tensorboard event file respectively.
If you'd like a different list of writers, you can overwrite it in
your trainer.
Returns:
list[EventWriter]: a list of :class:`EventWriter` objects.
It is now implemented by:
.. code-block:: python
return [
CommonMetricPrinter(self.max_iter),
JSONWriter(os.path.join(self.cfg.OUTPUT_DIR, "metrics.json")),
TensorboardXWriter(self.cfg.OUTPUT_DIR),
]
"""
# Assume the default print/log frequency.
# TODO: customize my writers
return [
# It may not always print what you want to see, since it prints "common" metrics only.
CommonMetricPrinter(self.max_iter),
JSONWriter(os.path.join(self.cfg.OUTPUT_DIR, "metrics.json")),
TensorboardXWriter(self.cfg.OUTPUT_DIR),
]
def _write_metrics(self, loss_dict: Dict[str, torch.Tensor],
data_time: float):
"""
Args:
loss_dict (dict): dict of scalar losses
data_time (float): time taken by the dataloader iteration
"""
device = next(iter(loss_dict.values())).device
# Use a new stream so these ops don't wait for DDP or backward
with torch.cuda.stream(torch.cuda.Stream() if device.type ==
"cuda" else None):
metrics_dict = {
k: v.detach().cpu().item()
for k, v in loss_dict.items()
}
metrics_dict["data_time"] = data_time
# Gather metrics among all workers for logging
# This assumes we do DDP-style training, which is currently the only
# supported method in detectron2.
all_metrics_dict = comm.gather(metrics_dict)
if comm.is_main_process():
storage = get_event_storage()
# data_time among workers can have high variance. The actual latency
# caused by data_time is the maximum among workers.
data_time = np.max([x.pop("data_time") for x in all_metrics_dict])
storage.put_scalar("data_time", data_time)
# average the rest metrics
metrics_dict = {
k: np.mean([x[k] for x in all_metrics_dict])
for k in all_metrics_dict[0].keys()
}
total_losses_reduced = sum(metrics_dict.values())
if not np.isfinite(total_losses_reduced):
raise FloatingPointError(
f"Loss became infinite or NaN at iteration={self.iter}!\n"
f"loss_dict = {metrics_dict}")
storage.put_scalar("total_loss", total_losses_reduced)
if len(metrics_dict) > 1:
storage.put_scalars(**metrics_dict)
class SpCL_UDA_Trainer(TrainerBase):
"""
load an un-pretrained model and train on the source & target domain from scratch
"""
def __init__(self, cfg):
super().__init__()
logger = logging.getLogger("fastreid")
if not logger.isEnabledFor(
logging.INFO): # if setup_logger is not called for fastreid
setup_logger()
# Create datasets
logger.info("==> Load source-domain dataset")
self.src = src = self.load_dataset(cfg.DATASETS.SRC)
self.src_pid_nums = src.get_num_pids(src.train)
logger.info("==> Load target-domain dataset")
self.tgt = tgt = self.load_dataset(cfg.DATASETS.TGT)
self.tgt_nums = len(tgt.train)
# Create model
self.model = self.build_model(cfg,
load_model=False,
show_model=False,
use_dsbn=True)
# Create hybrid memorys
self.hm = HybridMemory(num_features=cfg.MODEL.BACKBONE.FEAT_DIM,
num_samples=self.src_pid_nums + self.tgt_nums,
temp=cfg.MEMORY.TEMP,
momentum=cfg.MEMORY.MOMENTUM,
use_half=cfg.SOLVER.AMP.ENABLED).cuda()
# Initialize source-domain class centroids
logger.info(
"==> Initialize source-domain class centroids in the hybrid memory"
)
with inference_context(self.model), torch.no_grad():
src_train = self.build_dataset(cfg,
src.train,
is_train=False,
relabel=False,
with_mem_idx=False)
src_init_feat_loader = self.build_test_loader(cfg, src_train)
src_fname_feat_dict, _ = extract_features(self.model,
src_init_feat_loader)
src_feat_dict = collections.defaultdict(list)
for f, pid, _ in sorted(src.train):
src_feat_dict[pid].append(src_fname_feat_dict[f].unsqueeze(0))
src_centers = [
torch.cat(src_feat_dict[pid], 0).mean(0)
for pid in sorted(src_feat_dict.keys())
]
src_centers = torch.stack(src_centers, 0)
src_centers = F.normalize(src_centers, dim=1)
# Initialize target-domain instance features
logger.info(
"==> Initialize target-domain instance features in the hybrid memory"
)
with inference_context(self.model), torch.no_grad():
tgt_train = self.build_dataset(cfg,
tgt.train,
is_train=False,
relabel=False,
with_mem_idx=False)
tgt_init_feat_loader = self.build_test_loader(cfg, tgt_train)
tgt_fname_feat_dict, _ = extract_features(self.model,
tgt_init_feat_loader)
tgt_features = torch.cat([
tgt_fname_feat_dict[f].unsqueeze(0)
for f, _, _ in sorted(self.tgt.train)
], 0)
tgt_features = F.normalize(tgt_features, dim=1)
self.hm.features = torch.cat((src_centers, tgt_features), dim=0).cuda()
del (src_train, src_init_feat_loader, src_fname_feat_dict,
src_feat_dict, src_centers, tgt_train, tgt_init_feat_loader,
tgt_fname_feat_dict, tgt_features)
# Optimizer
self.optimizer, self.param_wrapper = self.build_optimizer(
cfg, self.model)
# For training, wrap with DDP. But don't need this for inference.
if comm.get_world_size() > 1:
# ref to https://github.com/pytorch/pytorch/issues/22049 to set `find_unused_parameters=True`
# for part of the parameters is not updated.
self.model = DistributedDataParallel(
self.model,
device_ids=[comm.get_local_rank()],
broadcast_buffers=False,
find_unused_parameters=True)
# Learning rate scheduler
self.iters_per_epoch = cfg.SOLVER.ITERS
self.scheduler = self.build_lr_scheduler(cfg, self.optimizer,
self.iters_per_epoch)
# Assume no other objects need to be checkpointed.
# We can later make it checkpoint the stateful hooks
self.checkpointer = Checkpointer(
# Assume you want to save checkpoints together with logs/statistics
self.model,
cfg.OUTPUT_DIR,
save_to_disk=comm.is_main_process(),
optimizer=self.optimizer,
**self.scheduler,
)
self.start_epoch = 0
self.max_epoch = cfg.SOLVER.MAX_EPOCH
self.max_iter = self.max_epoch * self.iters_per_epoch
self.warmup_iters = cfg.SOLVER.WARMUP_ITERS
self.delay_epochs = cfg.SOLVER.DELAY_EPOCHS
self.cfg = cfg
self.register_hooks(self.build_hooks())
if cfg.SOLVER.AMP.ENABLED:
unsupported = "AMPTrainer does not support single-process multi-device training!"
if isinstance(self.model, DistributedDataParallel):
assert not (self.model.device_ids
and len(self.model.device_ids) > 1), unsupported
from torch.cuda.amp.grad_scaler import GradScaler
self.grad_scaler = GradScaler()
else:
self.grad_scaler = None
def train(self):
"""
Run training.
Returns:
OrderedDict of results, if evaluation is enabled. Otherwise None.
"""
super().train(self.start_epoch, self.max_epoch, self.iters_per_epoch)
if comm.is_main_process():
assert hasattr(self, "_last_eval_results"
), "No evaluation results obtained during training!"
return self._last_eval_results
def before_train(self):
self.model.train()
if self.cfg.SOLVER.AMP.ENABLED:
assert torch.cuda.is_available(
), "CUDA is required for AMP training!"
return super().before_train()
def before_epoch(self):
logger = logging.getLogger('fastreid')
# Calculate distance
logger.info(
"==> Create pseudo labels for unlabeled target domain with self-paced policy"
)
tgt_features = self.hm.features[self.src_pid_nums:].clone()
rerank_dist = compute_jaccard_distance(tgt_features,
k1=self.cfg.CLUSTER.JACCARD.K1,
k2=self.cfg.CLUSTER.JACCARD.K2)
del tgt_features
if self.epoch == 0:
if self.cfg.CLUSTER.DBSCAN.ADAPTIVE_EPS:
logger.info("==> Calculating eps according to rerank_dist...")
tri_mat = np.triu(rerank_dist, 1) # tri_mat.dim=2
tri_mat = tri_mat[np.nonzero(tri_mat)] # tri_mat.dim=1
tri_mat = np.sort(tri_mat, axis=None)
top_num = np.round(self.cfg.SOLVER.RHO *
tri_mat.size).astype(int)
self.eps = tri_mat[:top_num].mean()
logger.info(f"==> epoch {self.epoch} eps: {self.eps}")
else:
self.eps = self.cfg.CLUSTER.DBSCAN.EPS
self.eps_tight = self.eps - self.cfg.CLUSTER.DBSCAN.EPS_GAP
self.eps_loose = self.eps + self.cfg.CLUSTER.DBSCAN.EPS_GAP
self.cluster = DBSCAN(eps=self.eps,
min_samples=4,
metric="precomputed",
n_jobs=-1)
self.cluster_tight = DBSCAN(eps=self.eps_tight,
min_samples=4,
metric="precomputed",
n_jobs=-1)
self.cluster_loose = DBSCAN(eps=self.eps_loose,
min_samples=4,
metric="precomputed",
n_jobs=-1)
# select & cluster images as training set of this epochs
pseudo_labels = self.cluster.fit_predict(rerank_dist)
pseudo_labels_tight = self.cluster_tight.fit_predict(rerank_dist)
pseudo_labels_loose = self.cluster_loose.fit_predict(rerank_dist)
num_ids = len(set(pseudo_labels)) - (1 if -1 in pseudo_labels else 0)
num_ids_tight = len(
set(pseudo_labels_tight)) - (1 if -1 in pseudo_labels_tight else 0)
num_ids_loose = len(
set(pseudo_labels_loose)) - (1 if -1 in pseudo_labels_loose else 0)
pseudo_labels = self.generate_pseudo_labels(pseudo_labels, num_ids)
pseudo_labels_tight = self.generate_pseudo_labels(
pseudo_labels_tight, num_ids_tight)
pseudo_labels_loose = self.generate_pseudo_labels(
pseudo_labels_loose, num_ids_loose)
# print(pseudo_labels.min(), pseudo_labels.max())
# exit()
# compute R_indep and R_comp
N = pseudo_labels.size(0)
label_sim = (pseudo_labels.expand(N, N).eq(
pseudo_labels.expand(N, N).t()).float()) # [N, N]
label_sim_tight = (pseudo_labels_tight.expand(N, N).eq(
pseudo_labels_tight.expand(N, N).t()).float())
label_sim_loose = (pseudo_labels_loose.expand(N, N).eq(
pseudo_labels_loose.expand(N, N).t()).float())
R_comp = 1 - torch.min(label_sim, label_sim_tight).sum(-1) / torch.max(
label_sim, label_sim_tight).sum(-1) # [N]
R_indep = 1 - torch.min(label_sim, label_sim_loose).sum(
-1) / torch.max(label_sim, label_sim_loose).sum(-1) # [N]
assert (R_comp.min() >= 0) and (R_comp.max() <= 1)
assert (R_indep.min() >= 0) and (R_indep.max() <= 1)
cluster_R_comp, cluster_R_indep = (
collections.defaultdict(list),
collections.defaultdict(list),
)
cluster_img_num = collections.defaultdict(int)
for i, (comp, indep,
label) in enumerate(zip(R_comp, R_indep, pseudo_labels)):
cluster_R_comp[label.item() - self.src_pid_nums].append(
comp.item())
cluster_R_indep[label.item() - self.src_pid_nums].append(
indep.item())
cluster_img_num[label.item() - self.src_pid_nums] += 1
cluster_R_comp = [
min(cluster_R_comp[i]) for i in sorted(cluster_R_comp.keys())
]
cluster_R_indep = [
min(cluster_R_indep[i]) for i in sorted(cluster_R_indep.keys())
]
cluster_R_indep_noins = [
iou for iou, num in zip(cluster_R_indep,
sorted(cluster_img_num.keys()))
if cluster_img_num[num] > 1
]
if self.epoch <= self.start_epoch:
"""
constant threshold α for identifying independent clusters is defined
by the top-90% Rindep before the first epoch and remains the same for all the training process
"""
logger.info("==> calculate independ before first epoch")
self.indep_thres = np.sort(cluster_R_indep_noins)[min(
len(cluster_R_indep_noins) - 1,
np.round(len(cluster_R_indep_noins) * 0.9).astype("int"),
)]
pseudo_labeled_dataset = []
outliers = 0
for i, ((fname, _, cid),
label) in enumerate(zip(sorted(self.tgt.train),
pseudo_labels)):
indep_score = cluster_R_indep[label.item() - self.src_pid_nums]
comp_score = R_comp[i]
if (indep_score <= self.indep_thres) and (
comp_score.item() <=
cluster_R_comp[label.item() - self.src_pid_nums]):
pseudo_labeled_dataset.append((fname, label.item(), cid))
else:
pseudo_label = self.src_pid_nums + len(
cluster_R_indep) + outliers
pseudo_labeled_dataset.append((fname, pseudo_label, cid))
pseudo_labels[i] = pseudo_label
outliers += 1
# statistics of clusters and un-clustered instances
index2label = collections.defaultdict(int)
for label in pseudo_labels:
index2label[label.item()] += 1
cluster_label = []
outlier_label = []
for k, v in index2label.items():
if v == 1:
outlier_label.append(k)
else:
cluster_label.append(k)
print(f'cluster_label', min(cluster_label), max(cluster_label),
len(cluster_label))
print(f'outlier label', min(outlier_label), max(outlier_label),
len(outlier_label))
index2label = np.fromiter(index2label.values(), dtype=float)
logger.info(
"==> Statistics for epoch {}: {} clusters, {} un-clustered instances, R_indep threshold is {}"
.format(
self.epoch,
(index2label > 1).sum(),
(index2label == 1).sum(),
1 - self.indep_thres,
))
self.hm.labels = torch.cat(
(torch.arange(self.src_pid_nums), pseudo_labels)).cuda()
src_train = self.build_dataset(
self.cfg,
self.src.train,
is_train=True,
relabel=True, # relabel?
# relabel=False,
with_mem_idx=True)
self.src_train_loader = self.build_train_loader(
self.cfg,
train_set=src_train,
sampler=RandomMultipleGallerySampler(
src_train.img_items,
self.cfg.SOLVER.IMS_PER_BATCH // comm.get_world_size(),
self.cfg.DATALOADER.NUM_INSTANCE,
with_mem_idx=True),
with_mem_idx=True)
# self.src_load_iter = iter(self.src_train_loader)
pseudo_tgt_train = self.build_dataset(
self.cfg,
pseudo_labeled_dataset,
is_train=True,
relabel=True, # relabel?
# relabel=False,
with_mem_idx=True)
self.pseudo_tgt_train_loader = self.build_train_loader(
self.cfg,
train_set=pseudo_tgt_train,
sampler=RandomMultipleGallerySampler(
pseudo_tgt_train.img_items,
self.cfg.SOLVER.IMS_PER_BATCH // comm.get_world_size(),
self.cfg.DATALOADER.NUM_INSTANCE,
with_mem_idx=True),
with_mem_idx=True)
# self.tgt_load_iter = iter(self.pseudo_tgt_train_loader)
return super().before_epoch()
# generate new dataset and calculate cluster centers
def generate_pseudo_labels(self, cluster_id, num):
labels = []
outliers = 0
for i, ((fname, _, cid),
id) in enumerate(zip(sorted(self.tgt.train), cluster_id)):
if id != -1:
labels.append(self.src_pid_nums + id)
else:
labels.append(self.src_pid_nums + num + outliers)
outliers += 1
return torch.Tensor(labels).long()
def run_step(self):
assert self.model.training, f"[{self.__class__.__name__}] model was changed to eval mode!"
if self.cfg.SOLVER.AMP.ENABLED:
assert torch.cuda.is_available(
), f"[{self.__class__.__name__}] CUDA is required for AMP training!"
from torch.cuda.amp.autocast_mode import autocast
start = time.perf_counter()
# load data
src_inputs = self.src_train_loader.next()
tgt_inputs = self.pseudo_tgt_train_loader.next()
# src_inputs = next(self.src_load_iter)
# tgt_inputs = next(self.tgt_load_iter)
def _parse_data(inputs):
# print(len(inputs))
# for i in range(len(inputs)):
# print(i, type(inputs[i]), inputs[i])
imgs, _, pids, _, indices = inputs
return imgs.cuda(), pids.cuda(), indices
# process inputs
s_inputs, s_targets, s_indices = _parse_data(src_inputs)
t_inputs, t_targets, t_indices = _parse_data(tgt_inputs)
# print('src', s_targets, s_indices)
# print('tgt', t_targets, t_indices)
# exit()
# arrange batch for domain-specific BNP
device_num = torch.cuda.device_count()
B, C, H, W = s_inputs.size()
def reshape(inputs):
return inputs.view(device_num, -1, C, H, W)
s_inputs, t_inputs = reshape(s_inputs), reshape(t_inputs)
inputs = torch.cat((s_inputs, t_inputs), 1).view(-1, C, H, W)
data_time = time.perf_counter() - start
def _forward():
outputs = self.model(inputs)
if isinstance(outputs, dict):
f_out = outputs['features']
else:
f_out = outputs
# de-arrange batch
f_out = f_out.view(device_num, -1, f_out.size(-1))
f_out_s, f_out_t = f_out.split(f_out.size(1) // 2, dim=1)
f_out_s, f_out_t = f_out_s.contiguous().view(
-1,
f_out.size(-1)), f_out_t.contiguous().view(-1, f_out.size(-1))
# compute loss with the hybrid memory
# with autocast(enabled=False):
loss_s = self.hm(f_out_s, s_targets)
loss_t = self.hm(f_out_t, t_indices + self.src_pid_nums)
loss_dict = {'loss_s': loss_s, 'loss_t': loss_t}
return loss_dict
if self.cfg.SOLVER.AMP.ENABLED:
with autocast():
loss_dict = _forward()
losses = sum(loss_dict.values())
self.optimizer.zero_grad()
self.grad_scaler.scale(losses).backward()
self._write_metrics(loss_dict, data_time)
self.grad_scaler.step(self.optimizer)
self.grad_scaler.update()
else:
loss_dict = _forward()
losses = sum(loss_dict.values())
self.optimizer.zero_grad()
losses.backward()
self._write_metrics(loss_dict, data_time)
self.optimizer.step()
if isinstance(self.param_wrapper, ContiguousParams):
self.param_wrapper.assert_buffer_is_valid()
@classmethod
def load_dataset(cls, name):
logger = logging.getLogger(__name__)
logger.info(f"Preparing {name}")
_root = os.getenv("FASTREID_DATASETS", "/root/datasets")
data = DATASET_REGISTRY.get(name)(root=_root)
if comm.is_main_process():
data.show_train()
return data
@classmethod
def build_dataset(cls,
cfg,
img_items,
is_train=False,
relabel=False,
transforms=None,
with_mem_idx=False):
if transforms is None:
transforms = build_transforms(cfg, is_train=is_train)
if with_mem_idx:
sorted_img_items = sorted(img_items)
for i in range(len(sorted_img_items)):
sorted_img_items[i] += (i, )
return InMemoryDataset(sorted_img_items, transforms, relabel)
else:
return CommDataset(img_items, transforms, relabel)
@classmethod
def build_train_loader(cls,
cfg,
train_set=None,
sampler=None,
with_mem_idx=False):
logger = logging.getLogger('fastreid')
logger.info("Prepare training loader")
total_batch_size = cfg.SOLVER.IMS_PER_BATCH
mini_batch_size = total_batch_size // comm.get_world_size()
if sampler is None:
num_instance = cfg.DATALOADER.NUM_INSTANCE
sampler_name = cfg.DATALOADER.SAMPLER_TRAIN
logger.info("Using training sampler {}".format(sampler_name))
if sampler_name == "TrainingSampler":
sampler = samplers.TrainingSampler(len(train_set))
elif sampler_name == "NaiveIdentitySampler":
sampler = samplers.NaiveIdentitySampler(
train_set.img_items, mini_batch_size, num_instance)
elif sampler_name == "BalancedIdentitySampler":
sampler = samplers.BalancedIdentitySampler(
train_set.img_items, mini_batch_size, num_instance)
elif sampler_name == "SetReWeightSampler":
set_weight = cfg.DATALOADER.SET_WEIGHT
sampler = samplers.SetReWeightSampler(train_set.img_items,
mini_batch_size,
num_instance, set_weight)
elif sampler_name == "ImbalancedDatasetSampler":
sampler = samplers.ImbalancedDatasetSampler(
train_set.img_items)
else:
raise ValueError(
"Unknown training sampler: {}".format(sampler_name))
iters = cfg.SOLVER.ITERS
num_workers = cfg.DATALOADER.NUM_WORKERS
batch_sampler = BatchSampler(sampler, mini_batch_size, True)
train_loader = IterLoader(
DataLoader(
Preprocessor(train_set, with_mem_idx),
num_workers=num_workers,
batch_sampler=batch_sampler,
pin_memory=True,
),
length=iters,
)
# train_loader = DataLoaderX(
# comm.get_local_rank(),
# dataset=Preprocessor(train_set, with_mem_idx),
# num_workers=num_workers,
# batch_sampler=batch_sampler,
# collate_fn=fast_batch_collator,
# pin_memory=True,
# )
return train_loader
@classmethod
def build_test_loader(cls, cfg, test_set):
logger = logging.getLogger('fastreid')
logger.info("Prepare testing loader")
# test_loader = DataLoader(
# # Preprocessor(test_set),
# test_set,
# batch_size=cfg.TEST.IMS_PER_BATCH,
# num_workers=cfg.DATALOADER.NUM_WORKERS,
# shuffle=False,
# pin_memory=True,
# )
test_batch_size = cfg.TEST.IMS_PER_BATCH
mini_batch_size = test_batch_size // comm.get_world_size()
num_workers = cfg.DATALOADER.NUM_WORKERS
data_sampler = samplers.InferenceSampler(len(test_set))
batch_sampler = BatchSampler(data_sampler, mini_batch_size, False)
test_loader = DataLoaderX(
comm.get_local_rank(),
dataset=test_set,
batch_sampler=batch_sampler,
num_workers=num_workers, # save some memory
collate_fn=fast_batch_collator,
pin_memory=True,
)
return test_loader
@classmethod
def build_evaluator(cls, cfg, dataset_name, output_dir=None):
data_loader, num_query = build_reid_test_loader(
cfg, dataset_name=dataset_name)
return data_loader, ReidEvaluator(cfg, num_query, output_dir)
@classmethod
def test(cls, cfg, model):
"""
Args:
cfg (CfgNode):
model (nn.Module):
Returns:
dict: a dict of result metrics
"""
logger = logging.getLogger('fastreid')
results = OrderedDict()
dataset_name = cfg.DATASETS.TGT
logger.info("Prepare testing set")
try:
data_loader, evaluator = cls.build_evaluator(cfg, dataset_name)
except NotImplementedError:
logger.warn(
"No evaluator found. implement its `build_evaluator` method.")
results[dataset_name] = {}
results_i = inference_on_dataset(model,
data_loader,
evaluator,
flip_test=cfg.TEST.FLIP.ENABLED)
results[dataset_name] = results_i
if comm.is_main_process():
assert isinstance(
results, dict
), "Evaluator must return a dict on the main process. Got {} instead.".format(
results)
logger.info("Evaluation results for {} in csv format:".format(
dataset_name))
results_i['dataset'] = dataset_name
print_csv_format(results_i)
# if len(results) == 1:
# results = list(results.values())[0]
return results
@classmethod
def build_model(cls,
cfg,
load_model=True,
show_model=True,
use_dsbn=False):
cfg = cfg.clone() # cfg can be modified by model
cfg.defrost()
cfg.MODEL.DEVICE = "cpu"
model = build_model(cfg)
logger = logging.getLogger('fastreid')
if load_model:
pretrain_path = cfg.MODEL.PRETRAIN_PATH
try:
state_dict = torch.load(
pretrain_path, map_location=torch.device("cpu"))['model']
for layer in cfg.MODEL.IGNORE_LAYERS:
if layer in state_dict.keys():
del state_dict[layer]
logger.info(f"Loading pretrained model from {pretrain_path}")
except FileNotFoundError as e:
logger.info(
f"{pretrain_path} is not found! Please check this path.")
raise e
except KeyError as e:
logger.info(
"State dict keys error! Please check the state dict.")
raise e
incompatible = model.load_state_dict(state_dict, strict=False)
if incompatible.missing_keys:
logger.info(
get_missing_parameters_message(incompatible.missing_keys))
if incompatible.unexpected_keys:
logger.info(
get_unexpected_parameters_message(
incompatible.unexpected_keys))
if use_dsbn:
logger.info("==> Convert BN to Domain Specific BN")
convert_dsbn(model)
if show_model:
logger.info("Model:\n{}".format(model))
model.to(torch.device("cuda"))
return model
@classmethod
def build_optimizer(cls, cfg, model):
"""
Returns:
torch.optim.Optimizer:
It now calls :func:`fastreid.solver.build_optimizer`.
Overwrite it if you'd like a different optimizer.
"""
return build_optimizer(cfg, model)
@classmethod
def build_lr_scheduler(cls, cfg, optimizer, iters_per_epoch):
"""
It now calls :func:`fastreid.solver.build_lr_scheduler`.
Overwrite it if you'd like a different scheduler.
"""
return build_lr_scheduler(cfg, optimizer, iters_per_epoch)
def build_hooks(self):
"""
Build a list of default hooks, including timing, evaluation,
checkpointing, lr scheduling, precise BN, writing events.
Returns:
list[HookBase]:
"""
logger = logging.getLogger(__name__)
cfg = self.cfg.clone()
cfg.defrost()
cfg.DATALOADER.NUM_WORKERS = 0 # save some memory and time for PreciseBN
cfg.DATASETS.NAMES = tuple([cfg.TEST.PRECISE_BN.DATASET
]) # set dataset name for PreciseBN
ret = [
hooks.IterationTimer(),
hooks.LRScheduler(self.optimizer, self.scheduler),
]
# if cfg.TEST.PRECISE_BN.ENABLED and hooks.get_bn_modules(self.model):
# logger.info("Prepare precise BN dataset")
# ret.append(hooks.PreciseBN(
# # Run at the same freq as (but before) evaluation.
# self.model,
# # Build a new data loader to not affect training
# self.build_train_loader(cfg),
# cfg.TEST.PRECISE_BN.NUM_ITER,
# ))
if len(cfg.MODEL.FREEZE_LAYERS) > 0 and cfg.SOLVER.FREEZE_ITERS > 0:
ret.append(
hooks.LayerFreeze(
self.model,
cfg.MODEL.FREEZE_LAYERS,
cfg.SOLVER.FREEZE_ITERS,
))
# Do PreciseBN before checkpointer, because it updates the model and need to
# be saved by checkpointer.
# This is not always the best: if checkpointing has a different frequency,
# some checkpoints may have more precise statistics than others.
def test_and_save_results():
self._last_eval_results = self.test(self.cfg, self.model)
return self._last_eval_results
# Do evaluation before checkpointer, because then if it fails,
# we can use the saved checkpoint to debug.
ret.append(hooks.EvalHook(cfg.TEST.EVAL_PERIOD, test_and_save_results))
if comm.is_main_process():
ret.append(
hooks.PeriodicCheckpointer(self.checkpointer,
cfg.SOLVER.CHECKPOINT_PERIOD))
# run writers in the end, so that evaluation metrics are written
ret.append(hooks.PeriodicWriter(self.build_writers(), 200))
return ret
def build_writers(self):
"""
Build a list of writers to be used. By default it contains
writers that write metrics to the screen,
a json file, and a tensorboard event file respectively.
If you'd like a different list of writers, you can overwrite it in
your trainer.
Returns:
list[EventWriter]: a list of :class:`EventWriter` objects.
It is now implemented by:
.. code-block:: python
return [
CommonMetricPrinter(self.max_iter),
JSONWriter(os.path.join(self.cfg.OUTPUT_DIR, "metrics.json")),
TensorboardXWriter(self.cfg.OUTPUT_DIR),
]
"""
# Assume the default print/log frequency.
# TODO: customize my writers
return [
# It may not always print what you want to see, since it prints "common" metrics only.
CommonMetricPrinter(self.max_iter),
JSONWriter(os.path.join(self.cfg.OUTPUT_DIR, "metrics.json")),
TensorboardXWriter(self.cfg.OUTPUT_DIR),
]
def _write_metrics(self, loss_dict: Dict[str, torch.Tensor],
data_time: float):
"""
Args:
loss_dict (dict): dict of scalar losses
data_time (float): time taken by the dataloader iteration
"""
device = next(iter(loss_dict.values())).device
# Use a new stream so these ops don't wait for DDP or backward
with torch.cuda.stream(torch.cuda.Stream() if device.type ==
"cuda" else None):
metrics_dict = {
k: v.detach().cpu().item()
for k, v in loss_dict.items()
}
metrics_dict["data_time"] = data_time
# Gather metrics among all workers for logging
# This assumes we do DDP-style training, which is currently the only
# supported method in detectron2.
all_metrics_dict = comm.gather(metrics_dict)
if comm.is_main_process():
storage = get_event_storage()
# data_time among workers can have high variance. The actual latency
# caused by data_time is the maximum among workers.
data_time = np.max([x.pop("data_time") for x in all_metrics_dict])
storage.put_scalar("data_time", data_time)
# average the rest metrics
metrics_dict = {
k: np.mean([x[k] for x in all_metrics_dict])
for k in all_metrics_dict[0].keys()
}
total_losses_reduced = sum(metrics_dict.values())
if not np.isfinite(total_losses_reduced):
raise FloatingPointError(
f"Loss became infinite or NaN at iteration={self.iter}!\n"
f"loss_dict = {metrics_dict}")
storage.put_scalar("total_loss", total_losses_reduced)
if len(metrics_dict) > 1:
storage.put_scalars(**metrics_dict)
def do_train(checkpoint=None,
direction='amr',
split_both_decoder=False,
fp16=False):
assert direction in ('amr', 'text', 'both')
model, tokenizer = instantiate_model_and_tokenizer(
config['model'],
checkpoint=checkpoint,
additional_tokens_smart_init=config['smart_init'],
dropout=config['dropout'],
attention_dropout=config['attention_dropout'],
from_pretrained=config['warm_start'],
init_reverse=split_both_decoder,
penman_linearization=config['penman_linearization'],
collapse_name_ops=config['collapse_name_ops'],
use_pointer_tokens=config['use_pointer_tokens'],
raw_graph=config.get('raw_graph', False))
print(model)
print(model.config)
if checkpoint is not None:
print(f'Checkpoint restored ({checkpoint})!')
if direction == 'both' and split_both_decoder:
params_dir_enc = list(model.model.encoder.parameters())
params_dir_enc_check = {id(p) for p in params_dir_enc}
params_dir_dec = set()
params_dir_dec |= {
p
for p in model.model.decoder.parameters()
if id(p) not in params_dir_enc_check
}
params_dir_dec |= {
p
for p in model.rev.model.decoder.parameters()
if id(p) not in params_dir_enc_check
}
params_dir_dec = list(params_dir_dec)
optimizer = RAdam([
{
'params': params_dir_enc,
'lr': config['learning_rate']
},
{
'params': params_dir_dec,
'lr': config['learning_rate'] * 2
},
],
weight_decay=config['weight_decay'])
else:
optimizer = RAdam(model.parameters(),
lr=config['learning_rate'],
weight_decay=config['weight_decay'])
if checkpoint is not None:
optimizer.load_state_dict(torch.load(checkpoint)['optimizer'])
if config['scheduler'] == 'cosine':
scheduler = transformers.get_cosine_schedule_with_warmup(
optimizer,
num_warmup_steps=config['warmup_steps'],
num_training_steps=config['training_steps'])
elif config['scheduler'] == 'constant':
scheduler = transformers.get_constant_schedule_with_warmup(
optimizer, num_warmup_steps=config['warmup_steps'])
else:
raise ValueError
scaler = GradScaler(enabled=fp16)
train_loader = instantiate_loader(
config['train'],
tokenizer,
batch_size=config['batch_size'],
evaluation=False,
use_recategorization=config['use_recategorization'],
remove_longer_than=config['remove_longer_than'],
remove_wiki=config['remove_wiki'],
dereify=config['dereify'],
)
dev_gold_path = ROOT / 'data/tmp/dev-gold.txt'
dev_pred_path = ROOT / 'data/tmp/dev-pred.txt'
dev_loader = instantiate_loader(
config['dev'],
tokenizer,
batch_size=config['batch_size'],
evaluation=True,
out=dev_gold_path,
use_recategorization=config['use_recategorization'],
remove_wiki=config['remove_wiki'],
dereify=config['dereify'],
)
if direction == 'amr':
def train_step(engine, batch):
model.train()
x, y, extra = batch
model.amr_mode = True
with autocast(enabled=fp16):
loss, *_ = model(**x, **y)
scaler.scale((loss / config['accum_steps'])).backward()
return loss.item()
@torch.no_grad()
def eval_step(engine, batch):
model.eval()
x, y, extra = batch
model.amr_mode = True
loss, *_ = model(**x, **y)
return loss.item()
elif direction == 'text':
def train_step(engine, batch):
model.train()
x, y, extra = batch
x, y = reverse_direction(x, y)
model.rev.amr_mode = False
with autocast(enabled=fp16):
loss, *_ = model.rev(**x, **y)
scaler.scale((loss / config['accum_steps'])).backward()
return loss.item()
@torch.no_grad()
def eval_step(engine, batch):
model.eval()
x, y, extra = batch
x, y = reverse_direction(x, y)
model.rev.amr_mode = False
loss, *_ = model(**x, **y)
return loss.item()
elif direction == 'both':
def train_step(engine, batch):
model.train()
x, y, extra = batch
model.amr_mode = True
with autocast(enabled=fp16):
loss1, *_ = model(**x, **y)
scaler.scale((loss1 / config['accum_steps'] * 0.5)).backward()
loss1 = loss1.item()
x, y = reverse_direction(x, y)
model.rev.amr_mode = False
with autocast(enabled=fp16):
loss2, *_ = model.rev(**x, **y)
scaler.scale((loss2 / config['accum_steps'] * 0.5)).backward()
return loss1, loss2.item()
@torch.no_grad()
def eval_step(engine, batch):
model.eval()
x, y, extra = batch
model.amr_mode = True
loss1, *_ = model(**x, **y)
x, y = reverse_direction(x, y)
model.rev.amr_mode = False
loss2, *_ = model.rev(**x, **y)
return loss1.item(), loss2.item()
else:
raise ValueError
trainer = Engine(train_step)
evaluator = Engine(eval_step)
@trainer.on(Events.STARTED)
def update(engine):
print('training started!')
@trainer.on(Events.EPOCH_COMPLETED)
@trainer.on(Events.ITERATION_COMPLETED(every=config['accum_steps']))
def update(engine):
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), config['grad_norm'])
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad()
scheduler.step()
@trainer.on(Events.EPOCH_COMPLETED)
def log_trn_loss(engine):
log_msg = f"training epoch: {engine.state.epoch}"
if direction in ('amr', 'both'):
log_msg += f" | loss_amr: {engine.state.metrics['trn_amr_loss']:.3f}"
if direction in ('text', 'both'):
log_msg += f" | loss_text: {engine.state.metrics['trn_text_loss']:.3f}"
print(log_msg)
@trainer.on(Events.EPOCH_COMPLETED)
def run_dev_eval(engine):
dev_loader.batch_size = config['batch_size']
dev_loader.device = next(model.parameters()).device
evaluator.run(dev_loader)
if not config['best_loss']:
if direction in ('amr', 'both'):
@evaluator.on(Events.EPOCH_COMPLETED)
def smatch_eval(engine):
device = next(model.parameters()).device
dev_loader.device = device
graphs = predict_amrs(
dev_loader,
model,
tokenizer,
restore_name_ops=config['collapse_name_ops'])
write_predictions(dev_pred_path, tokenizer, graphs)
try:
smatch = compute_smatch(dev_gold_path, dev_pred_path)
except:
smatch = 0.
engine.state.metrics['dev_smatch'] = smatch
if direction in ('text', 'both'):
@evaluator.on(Events.EPOCH_COMPLETED)
def smatch_eval(engine):
device = next(model.parameters()).device
dev_loader.device = device
pred_sentences = predict_sentences(
dev_loader,
model.rev,
tokenizer,
beam_size=config['beam_size'])
bleu = compute_bleu(dev_loader.dataset.sentences,
pred_sentences)
engine.state.metrics['dev_bleu'] = bleu.score
@evaluator.on(Events.EPOCH_COMPLETED)
def log_dev_loss(engine):
log_msg = f"dev epoch: {trainer.state.epoch}"
if direction in ('amr', 'both'):
log_msg += f" | loss_amr: {engine.state.metrics['dev_amr_loss']:.3f}"
if not config['best_loss']:
log_msg += f" | smatch: {engine.state.metrics['dev_smatch']:.3f}"
if direction in ('text', 'both'):
log_msg += f" | loss_text: {engine.state.metrics['dev_text_loss']:.3f}"
if not config['best_loss']:
log_msg += f" | bleu: {engine.state.metrics['dev_bleu']:.3f}"
print(log_msg)
if direction == 'amr':
RunningAverage(output_transform=lambda out: out).attach(
trainer, 'trn_amr_loss')
RunningAverage(output_transform=lambda out: out).attach(
evaluator, 'dev_amr_loss')
elif direction == 'text':
RunningAverage(output_transform=lambda out: out).attach(
trainer, 'trn_text_loss')
RunningAverage(output_transform=lambda out: out).attach(
evaluator, 'dev_text_loss')
elif direction == 'both':
RunningAverage(output_transform=lambda out: out[0]).attach(
trainer, 'trn_amr_loss')
RunningAverage(output_transform=lambda out: out[1]).attach(
trainer, 'trn_text_loss')
RunningAverage(output_transform=lambda out: out[0]).attach(
evaluator, 'dev_amr_loss')
RunningAverage(output_transform=lambda out: out[1]).attach(
evaluator, 'dev_text_loss')
if config['log_wandb']:
from ignite.contrib.handlers.wandb_logger import WandBLogger
wandb_logger = WandBLogger(init=False)
if direction == 'amr':
wandb_logger.attach_output_handler(
trainer,
event_name=Events.ITERATION_COMPLETED,
tag="iterations/trn_amr_loss",
output_transform=lambda loss: loss)
elif direction == 'text':
wandb_logger.attach_output_handler(
trainer,
event_name=Events.ITERATION_COMPLETED,
tag="iterations/trn_text_loss",
output_transform=lambda loss: loss)
if direction == 'both':
wandb_logger.attach_output_handler(
trainer,
event_name=Events.ITERATION_COMPLETED,
tag="iterations/trn_amr_loss",
output_transform=lambda loss: loss[0])
wandb_logger.attach_output_handler(
trainer,
event_name=Events.ITERATION_COMPLETED,
tag="iterations/trn_text_loss",
output_transform=lambda loss: loss[1])
if direction == 'amr':
metric_names_trn = ['trn_amr_loss']
metric_names_dev = ['dev_amr_loss']
if not config['best_loss']:
metric_names_dev.append('dev_smatch')
elif direction == 'text':
metric_names_trn = ['trn_text_loss']
metric_names_dev = ['dev_text_loss']
if not config['best_loss']:
metric_names_dev.append('dev_bleu')
elif direction == 'both':
metric_names_trn = ['trn_amr_loss', 'trn_text_loss']
metric_names_dev = ['dev_amr_loss', 'dev_smatch']
if not config['best_loss']:
metric_names_dev.extend(['dev_text_loss', 'dev_bleu'])
wandb_logger.attach_output_handler(
trainer,
event_name=Events.EPOCH_COMPLETED,
tag="epochs",
metric_names=metric_names_trn,
global_step_transform=lambda *_: trainer.state.iteration,
)
wandb_logger.attach_output_handler(
evaluator,
event_name=Events.EPOCH_COMPLETED,
tag="epochs",
metric_names=metric_names_dev,
global_step_transform=lambda *_: trainer.state.iteration,
)
@trainer.on(Events.ITERATION_COMPLETED)
def wandb_log_lr(engine):
wandb.log({'lr': scheduler.get_last_lr()[0]},
step=engine.state.iteration)
if config['save_checkpoints']:
if direction in ('amr', 'both'):
if config['best_loss']:
prefix = 'best-loss-amr'
score_function = lambda x: 1 / evaluator.state.metrics[
'dev_amr_loss']
else:
prefix = 'best-smatch'
score_function = lambda x: evaluator.state.metrics['dev_smatch'
]
else:
if config['best_loss']:
prefix = 'best-loss-text'
score_function = lambda x: 1 / evaluator.state.metrics[
'dev_amr_loss']
else:
prefix = 'best-bleu'
score_function = lambda x: evaluator.state.metrics['dev_bleu']
to_save = {'model': model, 'optimizer': optimizer}
if config['log_wandb']:
where_checkpoints = str(wandb_logger.run.dir)
else:
root = ROOT / 'runs'
try:
root.mkdir()
except:
pass
where_checkpoints = root / str(len(list(root.iterdir())))
try:
where_checkpoints.mkdir()
except:
pass
where_checkpoints = str(where_checkpoints)
print(where_checkpoints)
handler = ModelCheckpoint(
where_checkpoints,
prefix,
n_saved=1,
create_dir=True,
score_function=score_function,
global_step_transform=global_step_from_engine(trainer),
)
evaluator.add_event_handler(Events.EPOCH_COMPLETED, handler, to_save)
model.cuda()
device = next(model.parameters()).device
train_loader.device = device
trainer.run(train_loader, max_epochs=config['max_epochs'])
def train_step(
model: FlowModel,
config: TrainConfig,
action: ActionFn,
optimizer: optim.Optimizer,
batch_size: int,
scheduler: Any = None,
scaler: GradScaler = None,
pre_model: FlowModel = None,
dkl_factor: float = 1.,
xi: torch.Tensor = None,
):
"""Perform a single training step.
TODO: Add `torch.device` to arguments for DDP.
"""
t0 = time.time()
# layers, prior = model['layers'], model['prior']
optimizer.zero_grad()
loss_dkl = torch.tensor(0.0)
if torch.cuda.is_available():
loss_dkl = loss_dkl.cuda()
if pre_model is not None:
pre_xi = pre_model.prior.sample_n(batch_size)
x = qed.ft_flow(pre_model.layers, pre_xi)
xi = qed.ft_flow_inv(pre_model.layers, x)
# with torch.cuda.amp.autocast():
x, xi, logq = apply_flow_to_prior(model.prior,
model.layers,
xi=xi, batch_size=batch_size)
logp = (-1.) * action(x)
dkl = calc_dkl(logp, logq)
ess = calc_ess(logp, logq)
qi = qed.batch_charges(xi)
q = qed.batch_charges(x)
plaq = logp / (config.beta * config.volume)
dq = torch.sqrt((q - qi) ** 2)
loss_dkl = dkl_factor * dkl
if scaler is not None:
scaler.scale(loss_dkl).backward()
scaler.step(optimizer)
scaler.update()
else:
loss_dkl.backward()
optimizer.step()
if scheduler is not None:
scheduler.step(loss_dkl)
metrics = {
'dt': time.time() - t0,
'ess': grab(ess),
'logp': grab(logp),
'logq': grab(logq),
'loss_dkl': grab(loss_dkl),
'q': grab(q),
'dq': grab(dq),
'plaq': grab(plaq),
}
return metrics