def check_parity(amp: bool, manual_reduction: bool):
# The API should be the exact same in between the sharded and non-sharded variants, generic closure
def closure(model,
scaler,
input_tensor,
should_accumulate,
_manual_reduction=False):
accumulate_steps = 3 if should_accumulate else 1
model.zero_grad()
def step():
if scaler is not None:
with torch.cuda.amp.autocast():
loss = model(input_tensor).abs().sum()
scaler.scale(loss).backward()
else:
loss = model(input_tensor).abs().sum()
loss.backward()
with model.no_sync() if should_accumulate else suppress():
for _ in range(accumulate_steps - 1):
step()
if not _manual_reduction:
step()
else:
with model.no_sync():
step()
model.reduce()
# Any model works. Add one different buffer per rank
model = _get_mlp()
model.register_buffer("test_buffer", torch.ones((1)) * rank)
model.to(device)
# Make sure that the model starts with non-trainable, so that we check for the buckets to be
# properly reassigned when/if this changes
next(model.parameters()).requires_grad = False
sharded_optimizer = OSS(params=model.parameters(),
optim=torch.optim.SGD,
lr=1e-4,
momentum=0.99)
sharded_ddp_model = ShardedDataParallel(
module=model,
sharded_optimizer=sharded_optimizer,
broadcast_buffers=True,
reduce_buffer_size=reduce_buffer_size,
reduce_fp16=fp16_reduction,
)
ddp_model_single = copy.deepcopy(model)
ddp_optimizer = torch.optim.SGD(ddp_model_single.parameters(),
lr=1e-4,
momentum=0.99)
ddp_model = DDP(ddp_model_single,
device_ids=[rank],
broadcast_buffers=True,
find_unused_parameters=True)
if fp16_reduction:
from dist.algorithms.ddp_com_hooks.default_hooks import fp16_compress_hook
ddp_model.register_comm_hook(
state=None, hook=fp16_compress_hook) # type: ignore
ddp_scaler = TorchGradScaler() if amp else None
sharded_ddp_scaler = ShardedGradScaler() if amp else None
# The model should be synchronized in between the ranks at construction time, check that
check_same_model_params(sharded_ddp_model, ddp_model)
# Typical training loop, check that we get the exact same results as DDP
for i in range(NUMBER_BATCHS):
input_tensor = torch.rand((BATCH_SIZE, 2)).to(device)
def closure_ddp(input_tensor=input_tensor):
return closure(ddp_model, ddp_scaler, input_tensor,
grad_accumulation)
def closure_sharded(input_tensor=input_tensor):
return closure(
sharded_ddp_model,
sharded_ddp_scaler,
input_tensor,
grad_accumulation,
_manual_reduction=manual_reduction,
)
# Step/scale both
if ddp_scaler is not None:
_ = closure_ddp(input_tensor)
ddp_scaler.step(ddp_optimizer)
ddp_scaler.update()
else:
ddp_optimizer.step(closure=closure_ddp)
if sharded_ddp_scaler is not None:
_ = closure_sharded(input_tensor)
sharded_ddp_scaler.step(sharded_optimizer)
sharded_ddp_scaler.update()
else:
sharded_optimizer.step(closure=closure_sharded)
check_same_model_params(sharded_ddp_model, ddp_model,
f"Rank: {rank} - Step {i} broke")
# Flip the trainability of the first parameter back and forth
if i == 0 and change_train_graph:
next(sharded_ddp_model.parameters()).requires_grad = not next(
sharded_ddp_model.parameters()).requires_grad
next(ddp_model.parameters()).requires_grad = not next(
ddp_model.parameters()).requires_grad
check_same_model_params(
sharded_ddp_model, ddp_model,
f"Rank: {rank} - Trainability refresh {i} broke")
def _maybe_init_amp(self):
if self.fp16 and self.amp_grad_scaler is None and torch.cuda.is_available(
):
self.amp_grad_scaler = GradScaler()
def train(model,
dataloaders: dict,
criterion,
optimizer,
metrics,
scheduler,
reconstructor,
rundir: Union[str, bytes, os.PathLike],
stopper,
device: torch.device,
num_epochs: int = 1000,
steps_per_epoch: int = 1000,
steps_per_validation_epoch: int = 1000,
steps_per_test_epoch: int = 100,
early_stopping_begins: int = 0,
max_flow: float = 2.5,
dali: bool = False,
fp16: bool = False):
# check our inputs
assert (isinstance(model, nn.Module))
assert (isinstance(criterion, nn.Module))
assert (isinstance(optimizer, torch.optim.Optimizer))
scaler = None
if fp16:
scaler = GradScaler()
# loop over number of epochs!
for epoch in trange(0, num_epochs):
# if our learning rate scheduler plateaus when validation metric saturates, we have to pass our "key metric" for
# our validation set. Else, just step every epoch
if scheduler.name == 'plateau' and epoch > 0:
if hasattr(metrics, 'latest_key'):
if 'val' in list(metrics.latest_key.keys()):
scheduler.step(metrics.latest_key['val'])
elif epoch > 0:
scheduler.step()
# update the learning rate for this epoch
min_lr = utils.get_minimum_learning_rate(optimizer)
# store the learning rate for this epoch in our metrics file
# print('min lr: {}'.format(min_lr))
metrics.update_lr(min_lr)
# loop over our training set!
model, metrics, _ = loop_one_epoch(dataloaders['train'],
model,
criterion,
optimizer,
metrics,
reconstructor,
steps_per_epoch,
train_mode=True,
device=device,
dali=dali,
fp16=fp16,
scaler=scaler)
# evaluate on validation set
with torch.no_grad():
model, metrics, examples = loop_one_epoch(
dataloaders['val'],
model,
criterion,
optimizer,
metrics,
reconstructor,
steps_per_validation_epoch,
train_mode=False,
device=device,
max_flow=max_flow,
dali=dali,
fp16=fp16,
scaler=scaler)
# some training protocols do not have test sets, so just reuse validation set for testing inference speed
key = 'test' if 'test' in dataloaders.keys() else 'val'
loader = dataloaders[key]
# evaluate how fast inference takes, without loss calculation, which for some models can have a significant
# speed impact
metrics = speedtest(loader,
model,
metrics,
steps_per_test_epoch,
device=device,
dali=dali,
fp16=fp16)
# use our metrics file to output graphs for this epoch
viz.visualize_logger(metrics.fname, examples)
# save a checkpoint
utils.checkpoint(model, rundir, epoch)
# # update latest models file
# projects.write_latest_model(config['model'], config['flow_generator'], rundir, config)
# input the latest validation loss to the early stopper
if stopper.name == 'early':
should_stop, _ = stopper(metrics.latest_key['val'])
elif stopper.name == 'learning_rate':
should_stop = stopper(min_lr)
else:
# every epoch, increment stopper
should_stop = stopper()
if should_stop:
log.info('Stopping criterion reached!')
break
return model
def train_fn(train_loader, teacher_model, model, criterion, optimizer, epoch,
scheduler, device):
if CFG.device == 'GPU':
scaler = GradScaler()
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
scores = AverageMeter()
# switch to train mode
model.train()
start = end = time.time()
global_step = 0
for step, (images, images_annot, labels) in enumerate(train_loader):
# measure data loading time
data_time.update(time.time() - end)
with torch.no_grad():
teacher_features, _, _ = teacher_model(images_annot.to(device))
images = images.to(device)
labels = labels.to(device)
batch_size = labels.size(0)
if CFG.device == 'GPU':
with autocast():
features, _, y_preds = model(images)
loss = criterion(teacher_features, features, y_preds, labels)
# record loss
losses.update(loss.item(), batch_size)
if CFG.gradient_accumulation_steps > 1:
loss = loss / CFG.gradient_accumulation_steps
scaler.scale(loss).backward()
grad_norm = torch.nn.utils.clip_grad_norm_(
model.parameters(), CFG.max_grad_norm)
if (step + 1) % CFG.gradient_accumulation_steps == 0:
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad()
global_step += 1
elif CFG.device == 'TPU':
features, _, y_preds = model(images)
loss = criterion(teacher_features, features, y_preds, labels)
# record loss
losses.update(loss.item(), batch_size)
if CFG.gradient_accumulation_steps > 1:
loss = loss / CFG.gradient_accumulation_steps
loss.backward()
grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(),
CFG.max_grad_norm)
if (step + 1) % CFG.gradient_accumulation_steps == 0:
xm.optimizer_step(optimizer, barrier=True)
optimizer.zero_grad()
global_step += 1
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if CFG.device == 'GPU':
if step % CFG.print_freq == 0 or step == (len(train_loader) - 1):
print('Epoch: [{0}][{1}/{2}] '
'Data {data_time.val:.3f} ({data_time.avg:.3f}) '
'Elapsed {remain:s} '
'Loss: {loss.val:.4f}({loss.avg:.4f}) '
'Grad: {grad_norm:.4f} '
#'LR: {lr:.6f} '
.format(
epoch+1, step, len(train_loader), batch_time=batch_time,
data_time=data_time, loss=losses,
remain=timeSince(start, float(step+1)/len(train_loader)),
grad_norm=grad_norm,
#lr=scheduler.get_lr()[0],
))
elif CFG.device == 'TPU':
if step % CFG.print_freq == 0 or step == (len(train_loader) - 1):
xm.master_print('Epoch: [{0}][{1}/{2}] '
'Data {data_time.val:.3f} ({data_time.avg:.3f}) '
'Elapsed {remain:s} '
'Loss: {loss.val:.4f}({loss.avg:.4f}) '
'Grad: {grad_norm:.4f} '
#'LR: {lr:.6f} '
.format(
epoch+1, step, len(train_loader), batch_time=batch_time,
data_time=data_time, loss=losses,
remain=timeSince(start, float(step+1)/len(train_loader)),
grad_norm=grad_norm,
#lr=scheduler.get_lr()[0],
))
return losses.avg
class ModelTrainer(object):
def __init__(self, embed_model, optimizer, scheduler, mixedprec, **kwargs):
self.__model__ = embed_model
## Optimizer (e.g. Adam or SGD)
Optimizer = importlib.import_module(
'optimizer.' + optimizer).__getattribute__('Optimizer')
self.__optimizer__ = Optimizer(self.__model__.parameters(), **kwargs)
## Learning rate scheduler
Scheduler = importlib.import_module(
'scheduler.' + scheduler).__getattribute__('Scheduler')
self.__scheduler__, self.lr_step = Scheduler(self.__optimizer__,
**kwargs)
## For mixed precision training
self.scaler = GradScaler()
self.mixedprec = mixedprec
assert self.lr_step in ['epoch', 'iteration']
# ## ===== ===== ===== ===== ===== ===== ===== =====
# ## Train network
# ## ===== ===== ===== ===== ===== ===== ===== =====
def train_network(self, loader, verbose):
self.__model__.train()
stepsize = loader.batch_size
counter = 0
index = 0
loss = 0
top1 = 0 # EER or accuracy
tstart = time.time()
for data, label in loader:
data = data.transpose(1, 0)
## Reset gradients
self.__model__.zero_grad()
## Forward and backward passes
if self.mixedprec:
with autocast():
nloss, prec1 = self.__model__(data.cuda(), label.cuda())
self.scaler.scale(nloss).backward()
self.scaler.step(self.__optimizer__)
self.scaler.update()
else:
nloss, prec1 = self.__model__(data.cuda(), label.cuda())
nloss.backward()
self.__optimizer__.step()
loss += nloss.detach().cpu()
top1 += prec1.detach().cpu()
counter += 1
index += stepsize
telapsed = time.time() - tstart
tstart = time.time()
if verbose:
sys.stdout.write("\rProcessing (%d) " % (index))
sys.stdout.write(
"Loss %f TEER/TAcc %2.3f%% - %.2f Hz " %
(loss / counter, top1 / counter, stepsize / telapsed))
sys.stdout.flush()
if self.lr_step == 'iteration': self.__scheduler__.step()
if self.lr_step == 'epoch': self.__scheduler__.step()
sys.stdout.write("\n")
return (loss / counter, top1 / counter)
## ===== ===== ===== ===== ===== ===== ===== =====
## Evaluate from list
## ===== ===== ===== ===== ===== ===== ===== =====
def evaluateFromList(self,
test_list,
test_path,
nDataLoaderThread,
transform,
print_interval=100,
num_eval=10,
**kwargs):
self.__model__.eval()
feats = {}
tstart = time.time()
## Read all lines
with open(test_list) as f:
lines = f.readlines()
## Get a list of unique file names
files = sum([x.strip().split(',')[-2:] for x in lines], [])
setfiles = list(set(files))
setfiles.sort()
## Define test data loader
test_dataset = test_dataset_loader(setfiles,
test_path,
transform=transform,
num_eval=num_eval,
**kwargs)
test_loader = torch.utils.data.DataLoader(
test_dataset,
batch_size=1,
shuffle=False,
num_workers=nDataLoaderThread,
drop_last=False,
)
## Extract features for every image
for idx, data in enumerate(test_loader):
inp1 = data[0][0].cuda()
ref_feat = self.__model__(inp1).detach().cpu()
feats[data[1][0]] = ref_feat
telapsed = time.time() - tstart
if idx % print_interval == 0:
sys.stdout.write(
"\rReading %d of %d: %.2f Hz, embedding size %d" %
(idx, len(setfiles), idx / telapsed, ref_feat.size()[1]))
print('')
all_scores = []
all_labels = []
tstart = time.time()
## Read files and compute all scores
for idx, line in enumerate(lines):
data = line.strip().split(',')
ref_feat = feats[data[1]]
com_feat = feats[data[2]]
score = F.cosine_similarity(ref_feat, com_feat)
all_scores.append(score)
all_labels.append(int(data[0]))
if idx % print_interval == 0:
telapsed = time.time() - tstart
sys.stdout.write("\rComputing %d of %d: %.2f Hz" %
(idx, len(lines), idx / telapsed))
sys.stdout.flush()
print('')
return (all_scores, all_labels)
## ===== ===== ===== ===== ===== ===== ===== =====
## Save parameters
## ===== ===== ===== ===== ===== ===== ===== =====
def saveParameters(self, path):
torch.save(self.__model__.state_dict(), path)
## ===== ===== ===== ===== ===== ===== ===== =====
## Load parameters
## ===== ===== ===== ===== ===== ===== ===== =====
def loadParameters(self, path):
self_state = self.__model__.state_dict()
loaded_state = torch.load(path)
for name, param in loaded_state.items():
origname = name
if name not in self_state:
if name not in self_state:
print("%s is not in the model." % origname)
continue
if self_state[name].size() != loaded_state[origname].size():
print("Wrong parameter length: %s, model: %s, loaded: %s" %
(origname, self_state[name].size(),
loaded_state[origname].size()))
continue
self_state[name].copy_(param)
def main():
opt = parse_args()
opt.sever_name = gethostname()
# --- CUDA setting ---
os.environ["CUDA_VISIBLE_DEVICES"] = opt.gpu_ids
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# --- Random Seed setting ---
if opt.random_seed is None:
opt.random_seed = random.randint(1, 10000)
random.seed(opt.random_seed)
np.random.seed(opt.random_seed)
torch.manual_seed(opt.random_seed)
if torch.cuda.is_available():
torch.cuda.manual_seed_all(opt.random_seed)
cudnn.deterministic = True
# --- PATH setting ---
save_result_dir = Path(__file__).parent / "results" / opt.experiment_name
opt.save_model_dir = str(save_result_dir / "trained_models")
opt.save_log_path = str(save_result_dir / "train.log")
mkdirs(opt.save_model_dir)
# --- Prepare DataLoader ---
train_loader, valid_loader = get_train_val_loader(opt)
opt.src_vocab_size = train_loader.dataset.src_vocab_size
opt.tgt_vocab_size = train_loader.dataset.tgt_vocab_size
# --- Prepare Model ---
model = Transformer(
src_vocab_size=opt.src_vocab_size,
tgt_vocab_size=opt.tgt_vocab_size,
max_position_num=opt.max_position_num,
d_model=opt.d_model,
head_num=opt.head_num,
d_k=opt.d_k,
d_v=opt.d_v,
d_inner=opt.d_inner,
layer_num=opt.layer_num,
dropout=opt.dropout,
shared_embedding=opt.shared_embedding,
share_dec_input_output_embed=opt.share_dec_input_output_embed,
init_weight=opt.init_weight,
fused_layer_norm=opt.use_fused,
).to(device)
# --- Prepare optimizer and scaler ---
if opt.use_fused:
from apex.optimizers import FusedAdam as Adam
else:
from torch.optim import Adam
optimizer = Adam(filter(lambda x: x.requires_grad, model.parameters()),
betas=(0.9, 0.98),
eps=1e-09,
weight_decay=opt.weight_decay)
scaler = GradScaler(init_scale=65536.0, enabled=opt.use_amp)
# --- Restart setting ---
start_cnt = 1
steps_cnt = 0
if opt.adapt_NMT is not None:
ex_name, step_cnt = opt.adapt_NMT.split(',')
saved_path = f"{Path(__file__).parent}/results/{ex_name}/trained_models/step_{step_cnt}.pth"
saved_dict = torch.load(saved_path,
map_location=lambda storage, loc: storage)
check_arguments(saved_dict["settings"], opt)
model.load_state_dict(saved_dict["model"])
print(f"[Info]Loading complete ({saved_path})")
if opt.restart is not None:
start_cnt = opt.restart + 1
if opt.restart < 500:
model_name = f"epoch_{opt.restart}.pth"
else:
model_name = f"step_{opt.restart}.pth"
saved_path = f"{opt.save_model_dir}/{model_name}"
saved_dict = torch.load(saved_path,
map_location=lambda storage, loc: storage)
model.load_state_dict(saved_dict["model"])
optimizer.load_state_dict(saved_dict["optimizer"])
scaler.load_state_dict(saved_dict["scaler"])
steps_cnt = saved_dict["steps_cnt"]
print(f"[Info]Loading complete ({saved_path})")
scheduler = Scheduler(
optimizer=optimizer,
init_lr=0.,
end_lr=opt.end_lr,
warmup_steps=opt.warmup_steps,
current_steps=steps_cnt,
)
# --- DataParallel setting ---
gpus = [i for i in range(len(opt.gpu_ids.split(',')))]
if len(gpus) > 1:
model = nn.DataParallel(model, device_ids=gpus)
# --- Prepare trainer and validator ---
validator = ScoreCalculator(
model=model,
data_loader=valid_loader,
references=valid_loader.dataset.tgt_insts,
bpe=opt.bpe,
cp_avg_num=opt.check_point_average,
)
trainer = NMTTrainer(
model=model,
train_loader=train_loader,
optimizer=optimizer,
scaler=scaler,
scheduler=scheduler,
opt=opt,
validator=validator,
)
# --- Train ---
if opt.max_epoch is not None:
trainer.train_by_epoch(start_cnt)
else:
trainer.train_by_step(start_cnt)
class Trainer:
_MASTER_ADDR = 'localhost'
_MASTER_PORT = '12355'
def __init__(self, experiment_dir, train_dataset_dir, valid_dataset_dir,
gpt2_name_or_path, init_weights_from_checkpoint,
worker_batch_size, data_shuffle_seed, freeze_n_layers,
learning_rate, n_epochs, validate_each_n_steps, warmup_ratio,
n_accum_steps):
self._experiment_dir = Path(experiment_dir)
self._train_dataset_dir = Path(train_dataset_dir)
self._valid_dataset_dir = Path(valid_dataset_dir)
self._gpt2_name_or_path = gpt2_name_or_path
self._init_weights_from_checkpoint = init_weights_from_checkpoint
self._worker_batch_size = worker_batch_size
self._data_shuffle_seed = data_shuffle_seed
self._freeze_n_layers = freeze_n_layers
self._learning_rate = learning_rate
self._n_epochs = n_epochs
self._validate_each_n_steps = validate_each_n_steps
self._warmup_ratio = warmup_ratio
self._n_accum_steps = n_accum_steps or 1
self._world_size = torch.cuda.device_count()
self._tokenizer = load_tokenizer(self._train_dataset_dir)
self._optimizer = None
self._scaler = None
self._rank = None
self._model = None
self._train_dl = None
self._valid_dl = None
self._global_step = None
self._samples_seen = None
self._writer = None
self._model_params = None
checkpoint_dir = self._experiment_dir / CHECKPOINTS_DIR_NAME
checkpoint_dir.mkdir(exist_ok=True)
self._checkpoint_file_path = checkpoint_dir / 'last.ckpt'
dataset_meta = read_meta(train_dataset_dir)
with open(self._experiment_dir / 'meta.json', 'w') as file:
json.dump(dataset_meta, file, indent=2)
def run(self):
get_pretrained_gpt2_with_lm_head(self._gpt2_name_or_path)
mp.spawn(self._train, nprocs=self._world_size, join=True)
def _save_checkpoint(self):
checkpoint = {
'scaler_state_dict': self._scaler.state_dict(),
'model_state_dict': self._model.state_dict(),
'optimizer_state_dict': self._optimizer.state_dict(),
'scheduler_state_dict': self._scheduler.state_dict(),
'global_step': self._global_step,
'samples_seen': self._samples_seen,
'world_size': self._world_size,
'gpt2_config_dict': self._model.module.gpt2.config.to_dict(),
'model_params': self._model_params
}
torch.save(checkpoint, self._checkpoint_file_path)
def _load_checkpoint(self):
checkpoint = torch.load(self._checkpoint_file_path, map_location='cpu')
checkpoint_world_size = checkpoint['world_size']
if checkpoint_world_size != self._world_size:
raise ValueError(
f'Checkpoint world size {checkpoint_world_size} does not match with the current '
f'world size {self._world_size}.')
self._scaler.load_state_dict(checkpoint['scaler_state_dict'])
self._model.load_state_dict(checkpoint['model_state_dict'])
self._optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
self._scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
self._global_step = checkpoint['global_step']
self._samples_seen = checkpoint['samples_seen']
self._train_dl = self._get_dataloader(
is_train=True, samples_offset=self._samples_seen)
def _load_only_weights_from_checkpoint(self):
checkpoint = torch.load(self._init_weights_from_checkpoint,
map_location='cpu')
self._model.load_state_dict(checkpoint['model_state_dict'])
def _train(self, rank):
_seed_everything(self._data_shuffle_seed)
self._setup_ddp(rank)
self._rank = rank
self._scaler = GradScaler()
self._model = self._get_model(self._rank)
self._optimizer = AdamW(params=self._model.parameters(),
lr=self._learning_rate)
self._global_step = 0
self._samples_seen = 0
self._train_dl = self._get_dataloader(is_train=True, samples_offset=0)
self._valid_dl = self._get_dataloader(is_train=False, samples_offset=0)
steps_per_epoch = len(self._train_dl)
num_training_steps = steps_per_epoch * self._n_epochs
num_warmup_steps = self._warmup_ratio * num_training_steps
self._scheduler = get_linear_schedule_with_warmup(
optimizer=self._optimizer,
num_warmup_steps=num_warmup_steps,
num_training_steps=num_training_steps)
if self._checkpoint_file_path.is_file():
self._load_checkpoint()
elif self._init_weights_from_checkpoint:
self._load_only_weights_from_checkpoint()
while True:
if self._rank == 0:
self._writer = self._writer or SummaryWriter(
self._experiment_dir / 'tb_logs')
self._train_dl = tqdm.tqdm(self._train_dl,
desc='Train step',
total=num_training_steps,
position=1,
initial=self._global_step)
for i_step, model_input in enumerate(self._train_dl):
train_losses_dict = self._train_step(model_input)
if rank == 0:
self._train_dl.set_postfix({
'samples_seen':
self._samples_seen,
'epoch':
self._global_step / steps_per_epoch
})
self._write_tb_logs(train_losses_dict)
self._write_tb_logs({
'learning-rate':
self._optimizer.param_groups[0]['lr']
})
self._write_tb_logs(
{'max_seq_len': model_input.input_ids.size()[1]})
if self._rank == 0 and self._global_step % self._validate_each_n_steps == 0:
valid_loss = self._validate()
self._save_checkpoint()
self._write_tb_logs({'loss/valid': valid_loss})
if self._global_step >= num_training_steps:
break
dist.destroy_process_group()
def _write_tb_logs(self, values_dict):
for tag, val in values_dict.items():
self._writer.add_scalar(tag=tag,
scalar_value=val,
global_step=self._global_step)
def _train_step(self, model_input):
self._model.train()
with autocast():
model_output = self._model(model_input)
loss = model_output.loss / self._n_accum_steps
lm_loss = model_output.lm_loss / self._n_accum_steps
cls_loss = model_output.cls_loss / self._n_accum_steps
self._scaler.scale(loss).backward()
if self._global_step % self._n_accum_steps == 0:
self._scaler.step(self._optimizer)
self._scaler.update()
self._optimizer.zero_grad()
dist.all_reduce(loss)
dist.all_reduce(lm_loss)
dist.all_reduce(cls_loss)
loss = (loss.item() / self._world_size) * self._n_accum_steps
lm_loss = (lm_loss.item() / self._world_size) * self._n_accum_steps
cls_loss = (cls_loss.item() / self._world_size) * self._n_accum_steps
samples_seen = torch.tensor(len(model_input.input_ids),
device=self._rank)
dist.all_reduce(samples_seen)
self._samples_seen += samples_seen.item()
self._global_step += 1
self._scheduler.step()
losses_dict = {
'loss/train': loss,
'lm_loss/train': lm_loss,
'cls_loss/train': cls_loss
}
return losses_dict
def _setup_ddp(self, rank):
os.environ['MASTER_ADDR'] = 'localhost'
os.environ['MASTER_PORT'] = '12355'
dist.init_process_group("nccl", rank=rank, world_size=self._world_size)
def _get_model(self, rank):
model_params = {
'gpt2_name_or_path': self._gpt2_name_or_path,
'vocab_size': self._tokenizer.vocab_size,
'n_classes': 2,
'end_of_speaker_2_token_id':
self._tokenizer.end_of_speaker_2_token_id,
'cls_loss_weight': 0.25
}
model = DialogModel(**model_params)
model = model.to(rank)
model = DistributedDataParallel(model, device_ids=[rank])
self._model_params = model_params
return model
def _get_dataloader(self, is_train, samples_offset):
return get_dataloader(dataset_dir=self._train_dataset_dir
if is_train else self._valid_dataset_dir,
distractor_p=0.5,
batch_size=self._worker_batch_size,
num_workers=2,
sort_chunk_size=self._worker_batch_size * 10,
samples_offset=samples_offset,
data_shuffle_seed=self._data_shuffle_seed,
is_distributed=is_train,
pad_token_id=self._tokenizer.pad_token_id,
end_of_speaker_1_token_id=self._tokenizer.
end_of_speaker_1_token_id,
end_of_speaker_2_token_id=self._tokenizer.
end_of_speaker_2_token_id)
@torch.no_grad()
def _validate(self):
self._model.eval()
loss = 0
valid_dl = tqdm.tqdm(self._valid_dl,
desc='Valid step',
total=len(self._valid_dl),
position=2)
for model_input in valid_dl:
with autocast():
model_output = self._model(model_input)
loss_on_step = model_output.loss
loss += loss_on_step.item()
loss /= len(valid_dl)
return loss
class Trainer(BaseTrainer):
def __init__(self, config):
super(Trainer, self).__init__(config)
self.datamanager = DataManger(config["data"])
# model
self.model = Baseline(
num_classes=self.datamanager.datasource.get_num_classes("train")
)
# summary model
summary(
self.model,
input_size=(3, 256, 128),
batch_size=config["data"]["batch_size"],
device="cpu",
)
# losses
cfg_losses = config["losses"]
self.criterion = Softmax_Triplet_loss(
num_class=self.datamanager.datasource.get_num_classes("train"),
margin=cfg_losses["margin"],
epsilon=cfg_losses["epsilon"],
use_gpu=self.use_gpu,
)
self.center_loss = CenterLoss(
num_classes=self.datamanager.datasource.get_num_classes("train"),
feature_dim=2048,
use_gpu=self.use_gpu,
)
# optimizer
cfg_optimizer = config["optimizer"]
self.optimizer = torch.optim.Adam(
self.model.parameters(),
lr=cfg_optimizer["lr"],
weight_decay=cfg_optimizer["weight_decay"],
)
self.optimizer_centerloss = torch.optim.SGD(
self.center_loss.parameters(), lr=0.5
)
# learing rate scheduler
cfg_lr_scheduler = config["lr_scheduler"]
self.lr_scheduler = WarmupMultiStepLR(
self.optimizer,
milestones=cfg_lr_scheduler["steps"],
gamma=cfg_lr_scheduler["gamma"],
warmup_factor=cfg_lr_scheduler["factor"],
warmup_iters=cfg_lr_scheduler["iters"],
warmup_method=cfg_lr_scheduler["method"],
)
# track metric
self.train_metrics = MetricTracker("loss", "accuracy")
self.valid_metrics = MetricTracker("loss", "accuracy")
# save best accuracy for function _save_checkpoint
self.best_accuracy = None
# send model to device
self.model.to(self.device)
self.scaler = GradScaler()
# resume model from last checkpoint
if config["resume"] != "":
self._resume_checkpoint(config["resume"])
def train(self):
for epoch in range(self.start_epoch, self.epochs + 1):
result = self._train_epoch(epoch)
if self.lr_scheduler is not None:
self.lr_scheduler.step()
result = self._valid_epoch(epoch)
# add scalars to tensorboard
self.writer.add_scalars(
"Loss",
{
"Train": self.train_metrics.avg("loss"),
"Val": self.valid_metrics.avg("loss"),
},
global_step=epoch,
)
self.writer.add_scalars(
"Accuracy",
{
"Train": self.train_metrics.avg("accuracy"),
"Val": self.valid_metrics.avg("accuracy"),
},
global_step=epoch,
)
# logging result to console
log = {"epoch": epoch}
log.update(result)
for key, value in log.items():
self.logger.info(" {:15s}: {}".format(str(key), value))
# save model
if (
self.best_accuracy == None
or self.best_accuracy < self.valid_metrics.avg("accuracy")
):
self.best_accuracy = self.valid_metrics.avg("accuracy")
self._save_checkpoint(epoch, save_best=True)
else:
self._save_checkpoint(epoch, save_best=False)
# save logs
self._save_logs(epoch)
def _train_epoch(self, epoch):
"""Training step"""
self.model.train()
self.train_metrics.reset()
with tqdm(total=len(self.datamanager.get_dataloader("train"))) as epoch_pbar:
epoch_pbar.set_description(f"Epoch {epoch}")
for batch_idx, (data, labels, _) in enumerate(
self.datamanager.get_dataloader("train")
):
# push data to device
data, labels = data.to(self.device), labels.to(self.device)
# zero gradient
self.optimizer.zero_grad()
self.optimizer_centerloss.zero_grad()
with autocast():
# forward batch
score, feat = self.model(data)
# calculate loss and accuracy
loss = (
self.criterion(score, feat, labels)
+ self.center_loss(feat, labels) * self.config["losses"]["beta"]
)
_, preds = torch.max(score.data, dim=1)
# backward parameters
# loss.backward()
self.scaler.scale(loss).backward()
# backward parameters for center_loss
for param in self.center_loss.parameters():
param.grad.data *= 1.0 / self.config["losses"]["beta"]
# optimize
# self.optimizer.step()
self.scaler.step(self.optimizer)
self.optimizer_centerloss.step()
self.scaler.update()
# update loss and accuracy in MetricTracker
self.train_metrics.update("loss", loss.item())
self.train_metrics.update(
"accuracy",
torch.sum(preds == labels.data).double().item() / data.size(0),
)
# update process bar
epoch_pbar.set_postfix(
{
"train_loss": self.train_metrics.avg("loss"),
"train_acc": self.train_metrics.avg("accuracy"),
}
)
epoch_pbar.update(1)
return self.train_metrics.result()
def _valid_epoch(self, epoch):
"""Validation step"""
self.model.eval()
self.valid_metrics.reset()
with torch.no_grad():
with tqdm(total=len(self.datamanager.get_dataloader("val"))) as epoch_pbar:
epoch_pbar.set_description(f"Epoch {epoch}")
for batch_idx, (data, labels, _) in enumerate(
self.datamanager.get_dataloader("val")
):
# push data to device
data, labels = data.to(self.device), labels.to(self.device)
with autocast():
# forward batch
score, feat = self.model(data)
# calculate loss and accuracy
loss = (
self.criterion(score, feat, labels)
+ self.center_loss(feat, labels)
* self.config["losses"]["beta"]
)
_, preds = torch.max(score.data, dim=1)
# update loss and accuracy in MetricTracker
self.valid_metrics.update("loss", loss.item())
self.valid_metrics.update(
"accuracy",
torch.sum(preds == labels.data).double().item() / data.size(0),
)
# update process bar
epoch_pbar.set_postfix(
{
"val_loss": self.valid_metrics.avg("loss"),
"val_acc": self.valid_metrics.avg("accuracy"),
}
)
epoch_pbar.update(1)
return self.valid_metrics.result()
def _save_checkpoint(self, epoch, save_best=True):
"""save model to file"""
state = {
"epoch": epoch,
"state_dict": self.model.state_dict(),
"center_loss": self.center_loss.state_dict(),
"optimizer": self.optimizer.state_dict(),
"optimizer_centerloss": self.optimizer_centerloss.state_dict(),
"lr_scheduler": self.lr_scheduler.state_dict(),
"best_accuracy": self.best_accuracy,
}
filename = os.path.join(self.checkpoint_dir, "model_last.pth")
self.logger.info("Saving last model: model_last.pth ...")
torch.save(state, filename)
if save_best:
filename = os.path.join(self.checkpoint_dir, "model_best.pth")
self.logger.info("Saving current best: model_best.pth ...")
torch.save(state, filename)
def _resume_checkpoint(self, resume_path):
"""Load model from checkpoint"""
if not os.path.exists(resume_path):
raise FileExistsError("Resume path not exist!")
self.logger.info("Loading checkpoint: {} ...".format(resume_path))
checkpoint = torch.load(resume_path, map_location=self.map_location)
self.start_epoch = checkpoint["epoch"] + 1
self.model.load_state_dict(checkpoint["state_dict"])
self.center_loss.load_state_dict(checkpoint["center_loss"])
self.optimizer.load_state_dict(checkpoint["optimizer"])
self.optimizer_centerloss.load_state_dict(checkpoint["optimizer_centerloss"])
self.lr_scheduler.load_state_dict(checkpoint["lr_scheduler"])
self.best_accuracy = checkpoint["best_accuracy"]
self.logger.info(
"Checkpoint loaded. Resume training from epoch {}".format(self.start_epoch)
)
def _save_logs(self, epoch):
"""Save logs from google colab to google drive"""
if os.path.isdir(self.logs_dir_saved):
shutil.rmtree(self.logs_dir_saved)
destination = shutil.copytree(self.logs_dir, self.logs_dir_saved)
def __init__(self, config):
super(Trainer, self).__init__(config)
self.datamanager = DataManger(config["data"])
# model
self.model = Baseline(
num_classes=self.datamanager.datasource.get_num_classes("train")
)
# summary model
summary(
self.model,
input_size=(3, 256, 128),
batch_size=config["data"]["batch_size"],
device="cpu",
)
# losses
cfg_losses = config["losses"]
self.criterion = Softmax_Triplet_loss(
num_class=self.datamanager.datasource.get_num_classes("train"),
margin=cfg_losses["margin"],
epsilon=cfg_losses["epsilon"],
use_gpu=self.use_gpu,
)
self.center_loss = CenterLoss(
num_classes=self.datamanager.datasource.get_num_classes("train"),
feature_dim=2048,
use_gpu=self.use_gpu,
)
# optimizer
cfg_optimizer = config["optimizer"]
self.optimizer = torch.optim.Adam(
self.model.parameters(),
lr=cfg_optimizer["lr"],
weight_decay=cfg_optimizer["weight_decay"],
)
self.optimizer_centerloss = torch.optim.SGD(
self.center_loss.parameters(), lr=0.5
)
# learing rate scheduler
cfg_lr_scheduler = config["lr_scheduler"]
self.lr_scheduler = WarmupMultiStepLR(
self.optimizer,
milestones=cfg_lr_scheduler["steps"],
gamma=cfg_lr_scheduler["gamma"],
warmup_factor=cfg_lr_scheduler["factor"],
warmup_iters=cfg_lr_scheduler["iters"],
warmup_method=cfg_lr_scheduler["method"],
)
# track metric
self.train_metrics = MetricTracker("loss", "accuracy")
self.valid_metrics = MetricTracker("loss", "accuracy")
# save best accuracy for function _save_checkpoint
self.best_accuracy = None
# send model to device
self.model.to(self.device)
self.scaler = GradScaler()
# resume model from last checkpoint
if config["resume"] != "":
self._resume_checkpoint(config["resume"])
def __init__(self, opt):
super(inpaintModel, self).__init__(opt)
self.counter = 0
train_opt = opt['train']
# set if data should be normalized (-1,1) or not (0,1)
if self.is_train:
z_norm = opt['datasets']['train'].get('znorm', False)
# define networks and load pretrained models
self.netG = networks.define_G(opt).to(self.device) # G
if self.is_train:
self.netG.train()
if train_opt['gan_weight']:
self.netD = networks.define_D(opt).to(self.device) # D
self.netD.train()
self.load() # load G and D if needed
self.which_model_G = opt['network_G']['which_model_G']
# define losses, optimizer and scheduler
if self.is_train:
"""
Setup network cap
"""
# define if the generator will have a final capping mechanism in the output
self.outm = train_opt.get('finalcap', None)
"""
Setup batch augmentations
"""
self.mixup = train_opt.get('mixup', None)
if self.mixup:
#TODO: cutblur and cutout need model to be modified so LR and HR have the same dimensions (1x)
self.mixopts = train_opt.get(
'mixopts', ["blend", "rgb", "mixup", "cutmix", "cutmixup"
]) #, "cutout", "cutblur"]
self.mixprob = train_opt.get(
'mixprob', [1.0, 1.0, 1.0, 1.0, 1.0]) #, 1.0, 1.0]
self.mixalpha = train_opt.get(
'mixalpha', [0.6, 1.0, 1.2, 0.7, 0.7]) #, 0.001, 0.7]
self.aux_mixprob = train_opt.get('aux_mixprob', 1.0)
self.aux_mixalpha = train_opt.get('aux_mixalpha', 1.2)
self.mix_p = train_opt.get('mix_p', None)
"""
Setup frequency separation
"""
self.fs = train_opt.get('fs', None)
self.f_low = None
self.f_high = None
if self.fs:
lpf_type = train_opt.get('lpf_type', "average")
hpf_type = train_opt.get('hpf_type', "average")
self.f_low = FilterLow(filter_type=lpf_type).to(self.device)
self.f_high = FilterHigh(filter_type=hpf_type).to(self.device)
"""
Initialize losses
"""
#Initialize the losses with the opt parameters
# Generator losses:
self.generatorlosses = losses.GeneratorLoss(opt, self.device)
# TODO: show the configured losses names in logger
# print(self.generatorlosses.loss_list)
# Discriminator loss:
if train_opt['gan_type'] and train_opt['gan_weight']:
self.cri_gan = True
diffaug = train_opt.get('diffaug', None)
dapolicy = None
if diffaug: #TODO: this if should not be necessary
dapolicy = train_opt.get(
'dapolicy', 'color,translation,cutout') #original
self.adversarial = losses.Adversarial(train_opt=train_opt,
device=self.device,
diffaug=diffaug,
dapolicy=dapolicy)
# D_update_ratio and D_init_iters are for WGAN
self.D_update_ratio = train_opt.get('D_update_ratio', 1)
self.D_init_iters = train_opt.get('D_init_iters', 0)
else:
self.cri_gan = False
"""
Prepare optimizers
"""
self.optGstep = False
self.optDstep = False
if self.cri_gan:
self.optimizers, self.optimizer_G, self.optimizer_D = optimizers.get_optimizers(
self.cri_gan, self.netD, self.netG, train_opt, logger,
self.optimizers)
else:
self.optimizers, self.optimizer_G = optimizers.get_optimizers(
None, None, self.netG, train_opt, logger, self.optimizers)
self.optDstep = True
"""
Prepare schedulers
"""
self.schedulers = schedulers.get_schedulers(
optimizers=self.optimizers,
schedulers=self.schedulers,
train_opt=train_opt)
#Keep log in loss class instead?
self.log_dict = OrderedDict()
"""
Configure SWA
"""
#https://pytorch.org/blog/pytorch-1.6-now-includes-stochastic-weight-averaging/
self.swa = opt.get('use_swa', False)
if self.swa:
self.swa_start_iter = train_opt.get('swa_start_iter', 0)
# self.swa_start_epoch = train_opt.get('swa_start_epoch', None)
swa_lr = train_opt.get('swa_lr', 0.0001)
swa_anneal_epochs = train_opt.get('swa_anneal_epochs', 10)
swa_anneal_strategy = train_opt.get('swa_anneal_strategy',
'cos')
#TODO: Note: This could be done in resume_training() instead, to prevent creating
# the swa scheduler and model before they are needed
self.swa_scheduler, self.swa_model = swa.get_swa(
self.optimizer_G, self.netG, swa_lr, swa_anneal_epochs,
swa_anneal_strategy)
self.load_swa() #load swa from resume state
logger.info('SWA enabled. Starting on iter: {}, lr: {}'.format(
self.swa_start_iter, swa_lr))
"""
If using virtual batch
"""
batch_size = opt["datasets"]["train"]["batch_size"]
virtual_batch = opt["datasets"]["train"].get(
'virtual_batch_size', None)
self.virtual_batch = virtual_batch if virtual_batch \
>= batch_size else batch_size
self.accumulations = self.virtual_batch // batch_size
self.optimizer_G.zero_grad()
if self.cri_gan:
self.optimizer_D.zero_grad()
"""
Configure AMP
"""
self.amp = load_amp and opt.get('use_amp', False)
if self.amp:
self.cast = autocast
self.amp_scaler = GradScaler()
logger.info('AMP enabled')
else:
self.cast = nullcast
# print network
"""
def main(local_rank):
dist.init_process_group(backend='nccl', init_method='env://')
cfg.local_rank = local_rank
torch.cuda.set_device(local_rank)
cfg.rank = dist.get_rank()
cfg.world_size = dist.get_world_size()
print("world_size:", cfg.world_size)
trainset = MXFaceDataset(root_dir=cfg.rec, local_rank=local_rank)
train_sampler = torch.utils.data.distributed.DistributedSampler(
trainset, shuffle=True)
train_loader = DataLoaderX(local_rank=local_rank,
dataset=trainset,
batch_size=cfg.batch_size,
sampler=train_sampler,
num_workers=0,
pin_memory=True,
drop_last=False)
backbone = backbones.iefficientnet(False, num_features=cfg.embedding_size)
backbone.train()
# Broadcast init parameters
if cfg.preload:
print("Loading Backbone Checkpoint '{}'".format(cfg.preload))
loc = 'cuda:{}'.format(cfg.local_rank)
backbone.load_state_dict(torch.load(cfg.preload, map_location=loc))
backbone = backbone.to(local_rank)
backbone.train()
for ps in backbone.parameters():
dist.broadcast(ps, 0)
# DDP
backbone = torch.nn.parallel.DistributedDataParallel(
module=backbone, broadcast_buffers=False, device_ids=[cfg.local_rank])
backbone.train()
# Memory classifer
dist_sample_classifer = DistSampleClassifier(rank=dist.get_rank(),
local_rank=local_rank,
world_size=cfg.world_size)
if cfg.pf_preload:
file_model = os.path.join(cfg.pf_preload,
str(cfg.local_rank) + '_pf.pth')
print("Loading pf Checkpoint '{}'".format(file_model))
loc = 'cuda:{}'.format(cfg.local_rank)
dist_sample_classifer.load_state_dict(
torch.load(file_model, map_location=loc))
dist_sample_classifer = dist_sample_classifer.to(local_rank)
# Margin softmax
margin_softmax = MarginSoftmax(s=64.0, m=0.4)
# Optimizer for backbone and classifer
optimizer = SGD(
[
{
'params': backbone.parameters()
},
{
'params': dist_sample_classifer.parameters(),
'lr': cfg.lr
# 'params': dist_sample_classifer.parameters()
}
],
lr=cfg.lr,
momentum=cfg.momentum,
weight_decay=cfg.weight_decay,
rescale=cfg.world_size)
# Lr scheduler
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer=optimizer,
lr_lambda=cfg.lr_func)
n_epochs = cfg.num_epoch
start_epoch = 0
if local_rank == 0:
writer = SummaryWriter(log_dir='logs/shows')
#
total_step = int(
len(trainset) / cfg.batch_size / dist.get_world_size() * cfg.num_epoch)
if dist.get_rank() == 0:
print("Total Step is: %d" % total_step)
losses = AverageMeter()
global_step = 0
train_start = time.time()
if cfg.fp16:
scaler = GradScaler()
for epoch in range(start_epoch, n_epochs):
train_sampler.set_epoch(epoch)
for step, (img, label) in enumerate(train_loader):
# total_label, norm_weight = dist_sample_classifer.prepare(
# label, optimizer)
if cfg.fp16:
with autocast():
total_label, norm_weight = dist_sample_classifer.prepare(
label, optimizer)
features = F.normalize(backbone(img))
# Features all-gather
total_features = torch.zeros(features.size()[0] *
cfg.world_size,
cfg.embedding_size,
device=local_rank)
dist.all_gather(
list(total_features.chunk(cfg.world_size, dim=0)),
features.data)
total_features.requires_grad = True
# Calculate logits
logits = dist_sample_classifer(total_features, norm_weight)
logits = margin_softmax(logits, total_label)
with torch.no_grad():
max_fc = torch.max(logits, dim=1, keepdim=True)[0]
dist.all_reduce(max_fc, dist.ReduceOp.MAX)
# Calculate exp(logits) and all-reduce
logits_exp = torch.exp(logits - max_fc)
logits_sum_exp = logits_exp.sum(dim=1, keepdims=True)
dist.all_reduce(logits_sum_exp, dist.ReduceOp.SUM)
# Calculate prob
logits_exp.div_(logits_sum_exp)
# Get one-hot
grad = logits_exp
index = torch.where(total_label != -1)[0]
one_hot = torch.zeros(index.size()[0],
grad.size()[1],
device=grad.device)
one_hot.scatter_(1, total_label[index, None], 1)
# Calculate loss
loss = torch.zeros(grad.size()[0],
1,
device=grad.device)
loss[index] = grad[index].gather(
1, total_label[index, None])
dist.all_reduce(loss, dist.ReduceOp.SUM)
loss_v = loss.clamp_min_(1e-30).log_().mean() * (-1)
# Calculate grad
grad[index] -= one_hot
grad.div_(features.size()[0])
else:
total_label, norm_weight = dist_sample_classifer.prepare(
label, optimizer)
features = F.normalize(backbone(img))
# Features all-gather
total_features = torch.zeros(features.size()[0] *
cfg.world_size,
cfg.embedding_size,
device=local_rank)
dist.all_gather(
list(total_features.chunk(cfg.world_size, dim=0)),
features.data)
total_features.requires_grad = True
# Calculate logits
logits = dist_sample_classifer(total_features, norm_weight)
logits = margin_softmax(logits, total_label)
with torch.no_grad():
max_fc = torch.max(logits, dim=1, keepdim=True)[0]
dist.all_reduce(max_fc, dist.ReduceOp.MAX)
# Calculate exp(logits) and all-reduce
logits_exp = torch.exp(logits - max_fc)
logits_sum_exp = logits_exp.sum(dim=1, keepdims=True)
dist.all_reduce(logits_sum_exp, dist.ReduceOp.SUM)
# Calculate prob
logits_exp.div_(logits_sum_exp)
# Get one-hot
grad = logits_exp
index = torch.where(total_label != -1)[0]
one_hot = torch.zeros(index.size()[0],
grad.size()[1],
device=grad.device)
one_hot.scatter_(1, total_label[index, None], 1)
# Calculate loss
loss = torch.zeros(grad.size()[0], 1, device=grad.device)
loss[index] = grad[index].gather(1, total_label[index,
None])
dist.all_reduce(loss, dist.ReduceOp.SUM)
loss_v = loss.clamp_min_(1e-30).log_().mean() * (-1)
# Calculate grad
grad[index] -= one_hot
grad.div_(features.size()[0])
if cfg.fp16:
scaler.scale(logits).backward(grad)
else:
logits.backward(grad)
if total_features.grad is not None:
total_features.grad.detach_()
x_grad = torch.zeros_like(features)
# Feature gradient all-reduce
dist.reduce_scatter(
x_grad, list(total_features.grad.chunk(cfg.world_size, dim=0)))
x_grad.mul_(cfg.world_size)
# Backward backbone
if cfg.fp16:
scaler.scale(features).backward(x_grad)
scaler.step(optimizer)
else:
features.backward(x_grad)
optimizer.step()
# Update classifer
dist_sample_classifer.update()
optimizer.zero_grad()
# losses.update(loss_v, 1)
# if cfg.fp16:
# scaler.update()
losses.update(loss_v, 1)
if cfg.local_rank == 0 and step % 50 == 0:
time_now = (time.time() - train_start) / 3600
time_total = time_now / ((global_step + 1) / total_step)
time_for_end = time_total - time_now
writer.add_scalar('time_for_end', time_for_end, global_step)
writer.add_scalar('loss', loss_v, global_step)
writer.add_scalar('loss_epoch', loss_v, scheduler.get_lr()[0])
# writer.add_scalar('loss_lr', loss_v, scheduler.get_lr()[0])
print(
"Speed %d samples/sec Loss %.4f Epoch: %d Global Step: %d lr: %0.6f Required: %1.f hours"
% ((cfg.batch_size * global_step /
(time.time() - train_start) * cfg.world_size),
losses.avg, epoch, global_step, scheduler.get_lr()[0],
time_for_end))
losses.reset()
if cfg.fp16:
scaler.update()
global_step += 1
scheduler.step()
if dist.get_rank() == 0:
if not os.path.exists(cfg.output):
os.makedirs(cfg.output)
torch.save(backbone.module.state_dict(),
os.path.join(cfg.output,
str(epoch) + 'backbone.pth'))
#save pf
for gpu_index in range(cfg.world_size):
if not os.path.exists(cfg.output):
os.makedirs(cfg.output)
torch.save(
dist_sample_classifer.state_dict(),
os.path.join(cfg.output,
str(dist.get_rank()) + '_pf.pth'))
dist.destroy_process_group()
class inpaintModel(BaseModel):
def __init__(self, opt):
super(inpaintModel, self).__init__(opt)
self.counter = 0
train_opt = opt['train']
# set if data should be normalized (-1,1) or not (0,1)
if self.is_train:
z_norm = opt['datasets']['train'].get('znorm', False)
# define networks and load pretrained models
self.netG = networks.define_G(opt).to(self.device) # G
if self.is_train:
self.netG.train()
if train_opt['gan_weight']:
self.netD = networks.define_D(opt).to(self.device) # D
self.netD.train()
self.load() # load G and D if needed
self.which_model_G = opt['network_G']['which_model_G']
# define losses, optimizer and scheduler
if self.is_train:
"""
Setup network cap
"""
# define if the generator will have a final capping mechanism in the output
self.outm = train_opt.get('finalcap', None)
"""
Setup batch augmentations
"""
self.mixup = train_opt.get('mixup', None)
if self.mixup:
#TODO: cutblur and cutout need model to be modified so LR and HR have the same dimensions (1x)
self.mixopts = train_opt.get(
'mixopts', ["blend", "rgb", "mixup", "cutmix", "cutmixup"
]) #, "cutout", "cutblur"]
self.mixprob = train_opt.get(
'mixprob', [1.0, 1.0, 1.0, 1.0, 1.0]) #, 1.0, 1.0]
self.mixalpha = train_opt.get(
'mixalpha', [0.6, 1.0, 1.2, 0.7, 0.7]) #, 0.001, 0.7]
self.aux_mixprob = train_opt.get('aux_mixprob', 1.0)
self.aux_mixalpha = train_opt.get('aux_mixalpha', 1.2)
self.mix_p = train_opt.get('mix_p', None)
"""
Setup frequency separation
"""
self.fs = train_opt.get('fs', None)
self.f_low = None
self.f_high = None
if self.fs:
lpf_type = train_opt.get('lpf_type', "average")
hpf_type = train_opt.get('hpf_type', "average")
self.f_low = FilterLow(filter_type=lpf_type).to(self.device)
self.f_high = FilterHigh(filter_type=hpf_type).to(self.device)
"""
Initialize losses
"""
#Initialize the losses with the opt parameters
# Generator losses:
self.generatorlosses = losses.GeneratorLoss(opt, self.device)
# TODO: show the configured losses names in logger
# print(self.generatorlosses.loss_list)
# Discriminator loss:
if train_opt['gan_type'] and train_opt['gan_weight']:
self.cri_gan = True
diffaug = train_opt.get('diffaug', None)
dapolicy = None
if diffaug: #TODO: this if should not be necessary
dapolicy = train_opt.get(
'dapolicy', 'color,translation,cutout') #original
self.adversarial = losses.Adversarial(train_opt=train_opt,
device=self.device,
diffaug=diffaug,
dapolicy=dapolicy)
# D_update_ratio and D_init_iters are for WGAN
self.D_update_ratio = train_opt.get('D_update_ratio', 1)
self.D_init_iters = train_opt.get('D_init_iters', 0)
else:
self.cri_gan = False
"""
Prepare optimizers
"""
self.optGstep = False
self.optDstep = False
if self.cri_gan:
self.optimizers, self.optimizer_G, self.optimizer_D = optimizers.get_optimizers(
self.cri_gan, self.netD, self.netG, train_opt, logger,
self.optimizers)
else:
self.optimizers, self.optimizer_G = optimizers.get_optimizers(
None, None, self.netG, train_opt, logger, self.optimizers)
self.optDstep = True
"""
Prepare schedulers
"""
self.schedulers = schedulers.get_schedulers(
optimizers=self.optimizers,
schedulers=self.schedulers,
train_opt=train_opt)
#Keep log in loss class instead?
self.log_dict = OrderedDict()
"""
Configure SWA
"""
#https://pytorch.org/blog/pytorch-1.6-now-includes-stochastic-weight-averaging/
self.swa = opt.get('use_swa', False)
if self.swa:
self.swa_start_iter = train_opt.get('swa_start_iter', 0)
# self.swa_start_epoch = train_opt.get('swa_start_epoch', None)
swa_lr = train_opt.get('swa_lr', 0.0001)
swa_anneal_epochs = train_opt.get('swa_anneal_epochs', 10)
swa_anneal_strategy = train_opt.get('swa_anneal_strategy',
'cos')
#TODO: Note: This could be done in resume_training() instead, to prevent creating
# the swa scheduler and model before they are needed
self.swa_scheduler, self.swa_model = swa.get_swa(
self.optimizer_G, self.netG, swa_lr, swa_anneal_epochs,
swa_anneal_strategy)
self.load_swa() #load swa from resume state
logger.info('SWA enabled. Starting on iter: {}, lr: {}'.format(
self.swa_start_iter, swa_lr))
"""
If using virtual batch
"""
batch_size = opt["datasets"]["train"]["batch_size"]
virtual_batch = opt["datasets"]["train"].get(
'virtual_batch_size', None)
self.virtual_batch = virtual_batch if virtual_batch \
>= batch_size else batch_size
self.accumulations = self.virtual_batch // batch_size
self.optimizer_G.zero_grad()
if self.cri_gan:
self.optimizer_D.zero_grad()
"""
Configure AMP
"""
self.amp = load_amp and opt.get('use_amp', False)
if self.amp:
self.cast = autocast
self.amp_scaler = GradScaler()
logger.info('AMP enabled')
else:
self.cast = nullcast
# print network
"""
TODO:
Network summary? Make optional with parameter
could be an selector between traditional print_network() and summary()
"""
#self.print_network() #TODO
#https://github.com/Yukariin/DFNet/blob/master/data.py
def random_mask(self,
height=256,
width=256,
min_stroke=1,
max_stroke=4,
min_vertex=1,
max_vertex=12,
min_brush_width_divisor=16,
max_brush_width_divisor=10):
mask = np.ones((height, width))
min_brush_width = height // min_brush_width_divisor
max_brush_width = height // max_brush_width_divisor
max_angle = 2 * np.pi
num_stroke = np.random.randint(min_stroke, max_stroke + 1)
average_length = np.sqrt(height * height + width * width) / 8
for _ in range(num_stroke):
num_vertex = np.random.randint(min_vertex, max_vertex + 1)
start_x = np.random.randint(width)
start_y = np.random.randint(height)
for _ in range(num_vertex):
angle = np.random.uniform(max_angle)
length = np.clip(
np.random.normal(average_length, average_length // 2), 0,
2 * average_length)
brush_width = np.random.randint(min_brush_width,
max_brush_width + 1)
end_x = (start_x + length * np.sin(angle)).astype(np.int32)
end_y = (start_y + length * np.cos(angle)).astype(np.int32)
cv2.line(mask, (start_y, start_x), (end_y, end_x), 0.,
brush_width)
start_x, start_y = end_x, end_y
if np.random.random() < 0.5:
mask = np.fliplr(mask)
if np.random.random() < 0.5:
mask = np.flipud(mask)
return torch.from_numpy(
mask.reshape((1, ) + mask.shape).astype(np.float32)).unsqueeze(0)
def masking_images(self):
mask = self.random_mask(height=self.var_L.shape[2],
width=self.var_L.shape[3]).cuda()
for i in range(self.var_L.shape[0] - 1):
mask = torch.cat([
mask,
self.random_mask(height=self.var_L.shape[2],
width=self.var_L.shape[3]).cuda()
],
dim=0)
#self.var_L=self.var_L * mask
return self.var_L * mask, mask
def masking_images_with_invert(self):
mask = self.random_mask(height=self.var_L.shape[2],
width=self.var_L.shape[3]).cuda()
for i in range(self.var_L.shape[0] - 1):
mask = torch.cat([
mask,
self.random_mask(height=self.var_L.shape[2],
width=self.var_L.shape[3]).cuda()
],
dim=0)
#self.var_L=self.var_L * mask
return self.var_L * mask, self.var_L * (1 - mask), mask
def feed_data(self, data, need_HR=True):
# LR images
if self.which_model_G == 'EdgeConnect' or self.which_model_G == 'PRVS':
self.var_L = data['LR'].to(self.device)
self.canny_data = data['img_HR_canny'].to(self.device)
self.grayscale_data = data['img_HR_gray'].to(self.device)
#self.mask = data['green_mask'].to(self.device)
else:
self.var_L = data['LR'].to(self.device)
if need_HR: # train or val
# HR images
self.var_H = data['HR'].to(self.device)
# discriminator references
input_ref = data.get('ref', data['HR'])
self.var_ref = input_ref.to(self.device)
def feed_data_batch(self, data, need_HR=True):
# LR
self.var_L = data
def optimize_parameters(self, step):
# G
# freeze discriminator while generator is trained to prevent BP
if self.cri_gan:
for p in self.netD.parameters():
p.requires_grad = False
# batch (mixup) augmentations
aug = None
if self.mixup:
self.var_H, self.var_L, mask, aug = BatchAug(
self.var_H, self.var_L, self.mixopts, self.mixprob,
self.mixalpha, self.aux_mixprob, self.aux_mixalpha, self.mix_p)
if self.which_model_G == 'Pluralistic':
# pluralistic needs the inpainted area as an image and not only the cut-out
self.var_L, img_inverted, mask = self.masking_images_with_invert()
else:
self.var_L, mask = self.masking_images()
### Network forward, generate inpainted fake
with self.cast():
# normal
if self.which_model_G == 'AdaFill' or self.which_model_G == 'MEDFE' or self.which_model_G == 'RFR' or self.which_model_G == 'LBAM' or self.which_model_G == 'DMFN' or self.which_model_G == 'partial' or self.which_model_G == 'Adaptive' or self.which_model_G == 'DFNet' or self.which_model_G == 'RN':
self.fake_H = self.netG(self.var_L, mask)
# 2 rgb images
elif self.which_model_G == 'CRA' or self.which_model_G == 'pennet' or self.which_model_G == 'deepfillv1' or self.which_model_G == 'deepfillv2' or self.which_model_G == 'Global' or self.which_model_G == 'crfill' or self.which_model_G == 'DeepDFNet':
self.fake_H, self.other_img = self.netG(self.var_L, mask)
# special
elif self.which_model_G == 'Pluralistic':
self.fake_H, self.kl_rec, self.kl_g = self.netG(
self.var_L, img_inverted, mask)
save_image(self.fake_H, "self.fake_H_pluralistic.png")
elif self.which_model_G == 'EdgeConnect':
self.fake_H, self.other_img = self.netG(
self.var_L, self.canny_data, self.grayscale_data, mask)
elif self.which_model_G == 'FRRN':
self.fake_H, mid_x, mid_mask = self.netG(self.var_L, mask)
elif self.which_model_G == 'PRVS':
self.fake_H, _, edge_small, edge_big = self.netG(
self.var_L, mask, self.canny_data)
elif self.which_model_G == 'CSA':
#out_c, out_r, csa, csa_d
coarse_result, self.fake_H, csa, csa_d = self.netG(
self.var_L, mask)
elif self.which_model_G == 'atrous':
self.fake_H = self.netG(self.var_L)
else:
print("Selected model is not implemented.")
# Merge inpainted data with original data in masked region
self.fake_H = self.var_L * mask + self.fake_H * (1 - mask)
#save_image(self.fake_H, 'self_fake_H.png')
#/with self.cast():
#self.fake_H = self.netG(self.var_L, mask)
#self.counter += 1
#save_image(mask, str(self.counter)+'mask_train.png')
#save_image(self.fake_H, str(self.counter)+'fake_H_train.png')
# batch (mixup) augmentations
# cutout-ed pixels are discarded when calculating loss by masking removed pixels
if aug == "cutout":
self.fake_H, self.var_H = self.fake_H * mask, self.var_H * mask
l_g_total = 0
"""
Calculate and log losses
"""
loss_results = []
# training generator and discriminator
# update generator (on its own if only training generator or alternatively if training GAN)
if (self.cri_gan is not True) or (step % self.D_update_ratio == 0
and step > self.D_init_iters):
with self.cast(
): # Casts operations to mixed precision if enabled, else nullcontext
# regular losses
loss_results, self.log_dict = self.generatorlosses(
self.fake_H, self.var_H, self.log_dict, self.f_low)
# additional losses, in case a model does output more than a normal image
###############################
# deepfillv2 / global / crfill / CRA
if self.which_model_G == 'deepfillv2' or self.which_model_G == 'Global' or self.which_model_G == 'crfill' or self.which_model_G == 'CRA':
L1Loss = nn.L1Loss()
l1_stage1 = L1Loss(self.other_img, self.var_H)
self.log_dict.update(l1_stage1=l1_stage1)
loss_results.append(l1_stage1)
# edge-connect
if self.which_model_G == 'EdgeConnect':
L1Loss = nn.L1Loss()
l1_edge = L1Loss(self.other_img, self.canny_data)
self.log_dict.update(l1_edge=l1_edge)
loss_results.append(l1_edge)
###############################
# csa
if self.which_model_G == 'CSA':
#coarse_result, refine_result, csa, csa_d = g_model(masked, mask)
L1Loss = nn.L1Loss()
recon_loss = L1Loss(coarse_result, self.var_H) + L1Loss(
self.fake_H, self.var_H)
from models.modules.csa_loss import ConsistencyLoss
cons = ConsistencyLoss()
cons_loss = cons(csa, csa_d, self.var_H, mask)
self.log_dict.update(recon_loss=recon_loss)
loss_results.append(recon_loss)
self.log_dict.update(cons_loss=cons_loss)
loss_results.append(cons_loss)
###############################
# pluralistic (encoder kl loss)
if self.which_model_G == 'Pluralistic':
loss_kl_rec = self.kl_rec.mean()
loss_kl_g = self.kl_g.mean()
self.log_dict.update(loss_kl_rec=loss_kl_rec)
loss_results.append(loss_kl_rec)
self.log_dict.update(loss_kl_g=loss_kl_g)
loss_results.append(loss_kl_g)
###############################
# deepfillv1
if self.which_model_G == 'deepfillv1':
from models.modules.deepfillv1_loss import ReconLoss
ReconLoss_ = ReconLoss(1, 1, 1, 1)
reconstruction_loss = ReconLoss_(self.var_H,
self.other_img,
self.fake_H, mask)
self.log_dict.update(
reconstruction_loss=reconstruction_loss)
loss_results.append(reconstruction_loss)
###############################
# pennet
if self.which_model_G == 'pennet':
L1Loss = nn.L1Loss()
if self.other_img is not None:
pyramid_loss = 0
for _, f in enumerate(self.other_img):
pyramid_loss += L1Loss(
f,
torch.nn.functional.interpolate(
self.var_H,
size=f.size()[2:4],
mode='bilinear',
align_corners=True))
self.log_dict.update(pyramid_loss=pyramid_loss)
loss_results.append(pyramid_loss)
###############################
# FRRN
if self.which_model_G == 'FRRN':
L1Loss = nn.L1Loss()
# generator step loss
for idx in range(len(mid_x) - 1):
mid_l1_loss = L1Loss(mid_x[idx] * mid_mask[idx],
self.var_H * mid_mask[idx])
self.log_dict.update(mid_l1_loss=mid_l1_loss)
loss_results.append(mid_l1_loss)
###############################
# PRVS
if self.which_model_G == 'PRVS':
L1Loss = nn.L1Loss()
#from models.modules.PRVS_loss import edge_loss
#[edge_small, edge_big]
#adv_loss_0 = self.edge_loss(fake_edge[1], real_edge)
#dv_loss_1 = self.edge_loss(fake_edge[0], F.interpolate(real_edge, scale_factor = 0.5))
#adv_loss_0 = edge_loss(self, edge_big, self.canny_data, self.grayscale_data)
#adv_loss_1 = edge_loss(self, edge_small, torch.nn.functional.interpolate(self.canny_data, scale_factor = 0.5))
# l1 instead of discriminator loss
edge_big_l1 = L1Loss(edge_big, self.canny_data)
edge_small_l1 = L1Loss(
edge_small,
torch.nn.functional.interpolate(self.canny_data,
scale_factor=0.5))
self.log_dict.update(edge_big_l1=edge_big_l1)
loss_results.append(edge_big_l1)
self.log_dict.update(edge_small_l1=edge_small_l1)
loss_results.append(edge_small_l1)
###############################
#for key, value in self.log_dict.items():
# print(key, value)
l_g_total += sum(loss_results) / self.accumulations
if self.cri_gan:
# adversarial loss
l_g_gan = self.adversarial(
self.fake_H,
self.var_ref,
netD=self.netD,
stage='generator',
fsfilter=self.f_high) # (sr, hr)
self.log_dict['l_g_gan'] = l_g_gan.item()
l_g_total += l_g_gan / self.accumulations
#/with self.cast():
if self.amp:
# call backward() on scaled loss to create scaled gradients.
self.amp_scaler.scale(l_g_total).backward()
else:
l_g_total.backward()
# only step and clear gradient if virtual batch has completed
if (step + 1) % self.accumulations == 0:
if self.amp:
# unscale gradients of the optimizer's params, call
# optimizer.step() if no infs/NaNs in gradients, else, skipped
self.amp_scaler.step(self.optimizer_G)
# Update GradScaler scale for next iteration.
self.amp_scaler.update()
#TODO: remove. for debugging AMP
#print("AMP Scaler state dict: ", self.amp_scaler.state_dict())
else:
self.optimizer_G.step()
self.optimizer_G.zero_grad()
self.optGstep = True
if self.cri_gan:
# update discriminator
# unfreeze discriminator
for p in self.netD.parameters():
p.requires_grad = True
l_d_total = 0
with self.cast(
): # Casts operations to mixed precision if enabled, else nullcontext
l_d_total, gan_logs = self.adversarial(
self.fake_H,
self.var_ref,
netD=self.netD,
stage='discriminator',
fsfilter=self.f_high) # (sr, hr)
for g_log in gan_logs:
self.log_dict[g_log] = gan_logs[g_log]
l_d_total /= self.accumulations
#/with autocast():
if self.amp:
# call backward() on scaled loss to create scaled gradients.
self.amp_scaler.scale(l_d_total).backward()
else:
l_d_total.backward()
# only step and clear gradient if virtual batch has completed
if (step + 1) % self.accumulations == 0:
if self.amp:
# unscale gradients of the optimizer's params, call
# optimizer.step() if no infs/NaNs in gradients, else, skipped
self.amp_scaler.step(self.optimizer_D)
# Update GradScaler scale for next iteration.
self.amp_scaler.update()
else:
self.optimizer_D.step()
self.optimizer_D.zero_grad()
self.optDstep = True
def test(self, data):
"""
# generating random mask for validation
self.var_L, mask = self.masking_images()
if self.which_model_G == 'Pluralistic':
# pluralistic needs the inpainted area as an image and not only the cut-out
self.var_L, img_inverted, mask = self.masking_images_with_invert()
else:
self.var_L, mask = self.masking_images()
"""
self.mask = data['green_mask'].float().to(self.device).unsqueeze(0)
if self.which_model_G == 'Pluralistic':
img_inverted = self.var_L * (1 - self.mask)
self.var_L = self.var_L * self.mask
self.netG.eval()
with torch.no_grad():
if self.is_train:
# normal
if self.which_model_G == 'AdaFill' or self.which_model_G == 'MEDFE' or self.which_model_G == 'RFR' or self.which_model_G == 'LBAM' or self.which_model_G == 'DMFN' or self.which_model_G == 'partial' or self.which_model_G == 'Adaptive' or self.which_model_G == 'DFNet' or self.which_model_G == 'RN':
self.fake_H = self.netG(self.var_L, self.mask)
# 2 rgb images
elif self.which_model_G == 'CRA' or self.which_model_G == 'pennet' or self.which_model_G == 'deepfillv1' or self.which_model_G == 'deepfillv2' or self.which_model_G == 'Global' or self.which_model_G == 'crfill' or self.which_model_G == 'DeepDFNet':
self.fake_H, _ = self.netG(self.var_L, self.mask)
# special
elif self.which_model_G == 'Pluralistic':
self.fake_H, _, _ = self.netG(self.var_L, img_inverted,
self.mask)
elif self.which_model_G == 'EdgeConnect':
self.fake_H, _ = self.netG(self.var_L, self.canny_data,
self.grayscale_data, self.mask)
elif self.which_model_G == 'FRRN':
self.fake_H, _, _ = self.netG(self.var_L, self.mask)
elif self.which_model_G == 'PRVS':
self.fake_H, _, _, _ = self.netG(self.var_L, self.mask,
self.canny_data)
elif self.which_model_G == 'CSA':
_, self.fake_H, _, _ = self.netG(self.var_L, self.mask)
elif self.which_model_G == 'atrous':
self.fake_H = self.netG(self.var_L)
else:
print("Selected model is not implemented.")
# Merge inpainted data with original data in masked region
self.fake_H = self.var_L * self.mask + self.fake_H * (1 - self.mask)
self.netG.train()
def get_current_log(self):
return self.log_dict
def get_current_visuals(self, need_HR=True):
out_dict = OrderedDict()
out_dict['LR'] = self.var_L.detach()[0].float().cpu()
out_dict['SR'] = self.fake_H.detach()[0].float().cpu()
if need_HR:
out_dict['HR'] = self.var_H.detach()[0].float().cpu()
#TODO for PPON ?
#if get stages 1 and 2
#out_dict['SR_content'] = ...
#out_dict['SR_structure'] = ...
return out_dict
def get_current_visuals_batch(self, need_HR=True):
out_dict = OrderedDict()
out_dict['LR'] = self.var_L.detach().float().cpu()
out_dict['SR'] = self.fake_H.detach().float().cpu()
if need_HR:
out_dict['HR'] = self.var_H.detach().float().cpu()
#TODO for PPON ?
#if get stages 1 and 2
#out_dict['SR_content'] = ...
#out_dict['SR_structure'] = ...
return out_dict
def print_network(self):
# Generator
s, n = self.get_network_description(self.netG)
if isinstance(self.netG, nn.DataParallel):
net_struc_str = '{} - {}'.format(
self.netG.__class__.__name__,
self.netG.module.__class__.__name__)
else:
net_struc_str = '{}'.format(self.netG.__class__.__name__)
logger.info('Network G structure: {}, with parameters: {:,d}'.format(
net_struc_str, n))
logger.info(s)
if self.is_train:
# Discriminator
if self.cri_gan:
s, n = self.get_network_description(self.netD)
if isinstance(self.netD, nn.DataParallel):
net_struc_str = '{} - {}'.format(
self.netD.__class__.__name__,
self.netD.module.__class__.__name__)
else:
net_struc_str = '{}'.format(self.netD.__class__.__name__)
logger.info(
'Network D structure: {}, with parameters: {:,d}'.format(
net_struc_str, n))
logger.info(s)
#TODO: feature network is not being trained, is it necessary to visualize? Maybe just name?
# maybe show the generatorlosses instead?
'''
if self.generatorlosses.cri_fea: # F, Perceptual Network
#s, n = self.get_network_description(self.netF)
s, n = self.get_network_description(self.generatorlosses.netF) #TODO
#s, n = self.get_network_description(self.generatorlosses.loss_list.netF) #TODO
if isinstance(self.generatorlosses.netF, nn.DataParallel):
net_struc_str = '{} - {}'.format(self.generatorlosses.netF.__class__.__name__,
self.generatorlosses.netF.module.__class__.__name__)
else:
net_struc_str = '{}'.format(self.generatorlosses.netF.__class__.__name__)
logger.info('Network F structure: {}, with parameters: {:,d}'.format(net_struc_str, n))
logger.info(s)
'''
def load(self):
load_path_G = self.opt['path']['pretrain_model_G']
if load_path_G is not None:
logger.info('Loading pretrained model for G [{:s}] ...'.format(
load_path_G))
strict = self.opt['path'].get('strict', None)
self.load_network(load_path_G, self.netG, strict)
if self.opt['is_train'] and self.opt['train']['gan_weight']:
load_path_D = self.opt['path']['pretrain_model_D']
if self.opt['is_train'] and load_path_D is not None:
logger.info('Loading pretrained model for D [{:s}] ...'.format(
load_path_D))
self.load_network(load_path_D, self.netD)
def load_swa(self):
if self.opt['is_train'] and self.opt['use_swa']:
load_path_swaG = self.opt['path']['pretrain_model_swaG']
if self.opt['is_train'] and load_path_swaG is not None:
logger.info(
'Loading pretrained model for SWA G [{:s}] ...'.format(
load_path_swaG))
self.load_network(load_path_swaG, self.swa_model)
def save(self, iter_step, latest=None, loader=None):
self.save_network(self.netG, 'G', iter_step, latest)
if self.cri_gan:
self.save_network(self.netD, 'D', iter_step, latest)
if self.swa:
# when training with networks that use BN
# # Update bn statistics for the swa_model only at the end of training
# if not isinstance(iter_step, int): #TODO: not sure if it should be done only at the end
self.swa_model = self.swa_model.cpu()
torch.optim.swa_utils.update_bn(loader, self.swa_model)
self.swa_model = self.swa_model.cuda()
# Check swa BN statistics
# for module in self.swa_model.modules():
# if isinstance(module, torch.nn.modules.batchnorm._BatchNorm):
# print(module.running_mean)
# print(module.running_var)
# print(module.momentum)
# break
self.save_network(self.swa_model, 'swaG', iter_step, latest)
class BaseModel(ConfigModel):
def __init__(self, opt):
super().__init__()
self.opt = opt
self.nb_tasks = len(opt.num_classes)
self.gpu_ids = opt.gpu_ids
self.isTrain = opt.isTrain
self.save_dir = opt.checkpoints_dir # save all the checkpoints to save_dir
self.loss_names = []
self.net_names = []
self.visual_names = []
self.optimizers = []
self.image_paths = []
self.metric = 0 # used for learning rate policy 'plateau'
self.target, self.image = None, None
self.output = None
self.max_step = 0
@staticmethod
def modify_commandline_options(parser, is_train):
"""Add new model-specific options, and rewrite default values for existing options.
Parameters:
parser -- original option parser
is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.
Returns:
the modified parser.
"""
return parser
@staticmethod
def default_value(opt):
return opt
@exec_times(1)
def setup(self, task_index):
"""Load and print networks; create schedulers
Parameters:
opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
"""
self.task_index = task_index
logging.info("BaseModel set up")
if self.isTrain:
self.schedulers = [
get_scheduler(optimizer, self.opt)
for optimizer in self.optimizers
]
if not self.isTrain or self.opt.continue_train:
self.load_networks(self.opt.load_taskindex, self.opt.load_step,
self.opt.load_epoch)
if self.opt.amp:
self.scaler = GradScaler()
self.print_networks()
def eval(self):
"""Make models eval mode during test time"""
for name in self.net_names:
if isinstance(name, str):
net = getattr(self, "net_" + name)
net.eval()
def fit(self):
"""Make models train mode during train time"""
for name in self.net_names:
if isinstance(name, str):
net = getattr(self, "net_" + name)
net.train()
def train(self, task_index):
self.fit() # set train mode
self.optimizer.zero_grad(
set_to_none=True
) # clear network's existing gradients, torch.version>=1.7 is ok
# self.optimizer.zero_grad(set_grads_to_None=True) # clear network's existing gradients
# self.optimizer.zero_grad()
if self.opt.amp:
with autocast():
self.forward(
) # first call forward to calculate intermediate results
self.loss_total, self.losses_without_lambda = self.loss_criterion(
self.output, self.target, task_index)
else:
self.forward() # 多次执行得到的self.output会不一样？？？
self.loss_total, self.losses_without_lambda = self.loss_criterion(
self.output, self.target, task_index)
self.backward() # calculate gradients for network main
# self.scaler.unscale_(self.optimizer)
self.step()
if self.opt.amp:
self.scaler.update()
def step(self):
if self.opt.amp:
for optimizer in self.optimizers:
self.scaler.step(optimizer)
else:
for optimizer in self.optimizers:
optimizer.step() # update gradients for network
def test(self, visualizer=None):
"""Forward function used in test time.
This function wraps <forward> function in no_grad() so we don't save intermediate steps for backprop
It also calls <compute_visuals> to produce additional visualization results
"""
self.eval() # set eval mode
with torch.no_grad():
self.forward()
if visualizer is not None:
self.compute_visuals(visualizer)
def update_learning_rate(self):
"""Update learning rates for all the networks; called at the end of every epoch"""
for scheduler in self.schedulers:
if self.opt.lr_policy == 'plateau':
scheduler.step(self.metric)
else:
scheduler.step()
lr = self.optimizers[0].param_groups[0]['lr']
logging.info('learning rate = %.7f' % lr)
def get_current_losses(self):
"""Return traning losses_without_lambda / errors. fit.py will print out these errors on console, and save them to a file"""
errors_ret = OrderedDict()
for name in self.loss_names:
if isinstance(name, str):
errors_ret[name] = getattr(
self, name
) # float(...) works for both scalar tensor and float number
return errors_ret
def save_networks(self, task_index, step, epoch):
"""Save all the networks to the disk.
Parameters:
task_index (int) -- current trained task_index
step (int) -- current trained step
epoch (Union[int,str]) -- current epoch; used in the file tag '%s_%s_net_%s.pth' % (task_index,epoch, tag)
"""
task_index, step = int(task_index), int(step)
for file in os.listdir(self.save_dir):
if file.endswith('.pth'):
__task_index, __step, __epoch, *others = file.split('_')
__task_index, __step, __epoch = int(__task_index), int(
__step), int(__epoch)
if __task_index < task_index \
or __task_index == task_index and __step < step \
or __task_index == task_index and __step == step and (
__epoch < int(epoch) if epoch != 'best' else True):
path = os.path.join(self.save_dir, file)
logging.info(f'rmdir {path}')
# rm_dirs.append(path)
os.remove(path)
for name in self.net_names:
if isinstance(name, str):
save_filename = '%s_%s_%s_net_%s.pth' % (task_index, step,
epoch, name)
save_path = os.path.join(self.save_dir, save_filename)
net = getattr(self, "net_" + name)
if is_gpu_avaliable(self.opt):
torch.save(net.module.cpu().state_dict(), save_path)
net.to(device=self.opt.device)
else:
torch.save(net.cpu().state_dict(), save_path)
def load_networks(self, taks_index, step, epoch):
"""Load all the networks from the disk.
Parameters:
task_index (int) -- current trained task_index
epoch (int or str) -- current epoch or best epoch; used in the file tag '{taks_index}_{epoch}_net_{name}.pth' % (epoch, tag)
"""
for name in self.net_names:
if isinstance(name, str):
load_filename = f'{taks_index}_{step}_{epoch}_net_{name}.pth'
load_path = os.path.join(self.save_dir, load_filename)
if not os.path.exists(load_path):
logging.warning(
f'checkpoint path {load_path} is not exists, please checkout it!'
)
else:
net = getattr(self, "net_" + name)
if isinstance(net, torch.nn.DataParallel):
net = net.module
logging.info('loading the model from %s' % load_path)
state_dict = torch.load(load_path,
map_location=str(self.opt.device))
if hasattr(state_dict, '_metadata'):
del state_dict._metadata
# patch InstanceNorm checkpoints prior to 0.4
for key in list(state_dict.keys(
)): # need to copy keys here because we mutate in loop
self.__patch_instance_norm_state_dict(
state_dict, net, key.split('.'))
net.load_state_dict(state_dict)
def __patch_instance_norm_state_dict(self, state_dict, module, keys, i=0):
"""Fix InstanceNorm checkpoints incompatibility (prior to 0.4)"""
key = keys[i]
if i + 1 == len(keys): # at the end, pointing to a parameter/buffer
if module.__class__.__name__.startswith('InstanceNorm') and\
(key == 'running_mean' or key == 'running_var'):
if getattr(module, key) is None:
state_dict.pop('.'.join(keys))
if module.__class__.__name__.startswith('InstanceNorm') and\
(key == 'num_batches_tracked'):
state_dict.pop('.'.join(keys))
else:
self.__patch_instance_norm_state_dict(state_dict,
getattr(module, key), keys,
i + 1)
def print_networks(self):
"""Print the total number of parameters in the network and (if verbose) network architecture
Parameters:
verbose (bool) -- if verbose: print the network architecture
"""
logging.info('---------- Networks initialized -------------')
for name in self.net_names:
if isinstance(name, str):
net = getattr(self, "net_" + name)
num_params = 0
for param in net.parameters():
num_params += param.numel()
logging.info(net)
logging.info(
'[Network %s] Total number of parameters : %.3f M' %
(name, num_params / 1e6))
logging.info('-----------------------------------------------')
def get_matrix_item(self, task_index) -> MatrixItem:
"""for current task_index, after calling set_data() and forward(), then call <get_matrix_item()> will get MatrixItem """
if self.output is None or self.target is None:
raise ValueError(
f"Expected model.set_data(data) and model.forward() be called before model.get_matrix_item(), but not called {'forward()' if self.output is None else ''} {'and set_data()' if self.target is None else ''}"
)
return MatrixItem(self.output[task_index], self.target,
self.loss_criterion)
def init_net_with_dataparaller(self, opt, net):
"""Initialize a network: 1. register CPU/GPU device (with multi-GPU support); 2. initialize the network weights
Parameters:
net (network) -- the network to be initialized
init_type (str) -- the tag of an initialization method: normal | xavier | kaiming | orthogonal
gain (float) -- scaling factor for normal, xavier and orthogonal.
gpu_ids (int list) -- which GPUs the network runs on: e.rg., 0,1,2
Return an initialized network.
"""
if len(opt.gpu_ids) > 0:
assert (torch.cuda.is_available())
# for none distributed data parallel, only DataParallel
# net = DataParallel(net, device_ids=opt.gpu_ids)
# net = DataParallel(net, device_ids=[0])
net = DataParallel(net)
# self.cuda()
# net.to(device=self.opt.device)
logging.info(f'Data Paraller: DataParallel')
return net
else:
raise ValueError(f'Expected GPU training, but got none gpu_ids')
def __getattr__(self, item):
if "loss" in item and hasattr(self.loss_criterion, item):
return getattr(self.loss_criterion, item)
for net in self._get_all_nets():
if "net" in item and hasattr(net, item):
return getattr(net, item)
raise AttributeError(
f"Model {self.opt.model_name} object has no attribute '{item}'")
def set_data(self, data: PseudoData):
"""Unpack input _data from the dataloader and perform necessary pre-processing steps.
"""
self.image = data.image
self.target: MultiOutput = data.target
def forward(self):
"""Run forward pass. This will be called by both functions <optimize_parameters> and <test>."""
if self.image is None:
raise ValueError(
f"Expected model.set_data(data) be called before forward(), to get data"
)
task_outputs = self.net_main(self.image)
self.output: MultiOutput = MultiOutput(task_outputs)
# @torchsnooper.snoop()
def backward(self):
"""Calculate losses_without_lambda, gradients, and update network weights; called in every training iteration"""
if self.opt.amp:
for loss in self.losses_without_lambda[:-1]:
self.scaler.scale(loss).backward(retain_graph=True)
self.scaler.scale(self.losses_without_lambda[-1]).backward()
else:
for loss in self.losses_without_lambda[:-1]:
loss.backward(retain_graph=True)
self.losses_without_lambda[-1].backward()
def _get_all_nets(self):
nets = self.__dict__
return [nets["net_" + net_name] for net_name in self.net_names]
def compute_visuals(self, visualizer):
"""Calculate additional visualization"""
pass
def cuda(self, device=None):
"""cuda: net"""
if device is None:
device = self.opt.device
for name in self.net_names:
if isinstance(name, str):
net = getattr(self, "net_" + name)
# net=net.cuda(device)
net.to(device)
logging.info(f'[Network {name}] move to cuda({device})')
def init_optimizers(self, opt):
self.cuda()
if self.isTrain:
if self.opt.optimizer_type == 'adam':
self.optimizer = torch.optim.Adam(self.net_main.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
elif opt.optimizer_type == 'sgd':
self.optimizer = torch.optim.SGD(self.net_main.parameters(),
lr=opt.lr)
elif opt.optimizer_type == 'adamw':
self.optimizer = torch.optim.AdamW(self.net_main.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
else:
raise ValueError(
f"Expected opt.optimizer_type in ['adam','sgd'], but got {opt.optimizer_type}"
)
self.net_main = self.init_net_with_dataparaller(opt, self.net_main)
self.optimizers = [self.optimizer]
def main_worker(gpu, ngpus_per_node, args):
global global_step
global start_time
args.gpu = gpu
if args.gpu is not None:
print("Use GPU: {} for training".format(args.gpu))
if args.distributed:
if args.dist_url == "env://" and args.rank == -1:
args.rank = int(os.environ["RANK"])
if args.distributed:
args.rank = args.rank * ngpus_per_node + gpu
dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url,
world_size=args.world_size, rank=args.rank)
hps = Hyperparameters(args)
sample_path, save_path, load_path, log_path = mkdir(args)
if not args.distributed or (args.rank % ngpus_per_node == 0):
log, log_train, log_eval = get_logger(log_path, args.model_name)
else:
log, log_train, log_eval = None, None, None
model = build_model(hps, log)
if args.distributed: # Multiple processes, single GPU per process
if args.gpu is not None:
def _transform_(m):
return nn.parallel.DistributedDataParallel(
m, device_ids=[args.gpu], output_device=args.gpu, check_reduction=True)
torch.cuda.set_device(args.gpu)
model.cuda(args.gpu)
model.multi_gpu_wrapper(_transform_)
args.bsz = int(args.bsz / ngpus_per_node)
args.workers = 0
else:
assert 0, "DistributedDataParallel constructor should always set the single device scope"
elif args.gpu is not None: # Single process, single GPU per process
torch.cuda.set_device(args.gpu)
model = model.cuda(args.gpu)
else: # Single process, multiple GPUs per process
def _transform_(m):
return nn.DataParallel(m)
model = model.cuda()
model.multi_gpu_wrapper(_transform_)
train_loader, test_loader, synth_loader = load_dataset(args)
optimizer = optim.Adam(model.parameters(), lr=args.lr)
scaler = GradScaler()
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=args.step_size, gamma=args.gamma)
state = {k: v for k, v in args._get_kwargs()}
if args.load_step == 0:
# new model
global_epoch = 0
global_step = 0
actnorm_init(train_loader, model, args.gpu)
else:
# saved model
model, optimizer, scheduler, global_epoch, global_step = load_checkpoint(args.load_step, load_path, model, optimizer, scheduler)
if log is not None:
log.write('\n ! --- load the model and continue training --- ! \n')
log_train.write('\n ! --- load the model and continue training --- ! \n')
log_eval.write('\n ! --- load the model and continue training --- ! \n')
log.flush()
log_train.flush()
log_eval.flush()
start_time = time.time()
dateTime = datetime.datetime.fromtimestamp(start_time).strftime('%Y-%m-%d %H:%M:%S')
print('training starts at ', dateTime)
for epoch in range(global_epoch + 1, args.epochs + 1):
training_epoch_loss = train(args.gpu, epoch, train_loader, synth_loader, sample_path, model, optimizer, scaler, scheduler, log_train, args)
with torch.no_grad():
eval_epoch_loss = evaluate(args.gpu, epoch, test_loader, model, log_eval)
if log is not None:
state['training_loss'] = training_epoch_loss
state['eval_loss'] = eval_epoch_loss
state['epoch'] = epoch
log.write('%s\n' % json.dumps(state))
log.flush()
if not args.distributed or (args.rank % ngpus_per_node == 0):
save_checkpoint(save_path, model, optimizer, scaler, scheduler, global_step, epoch)
print('Epoch {} Model Saved! Loss : {:.4f}'.format(epoch, eval_epoch_loss))
with torch.no_grad():
synthesize(args.gpu, sample_path, synth_loader, model, args.num_sample, args.sr)
gc.collect()
if log is not None:
log_train.close()
log_eval.close()
log.close()
class ModelTrainer(object):
def __init__(self, speaker_model, optimizer, scheduler, gpu, mixedprec,
**kwargs):
self.__model__ = speaker_model
Optimizer = importlib.import_module(
'optimizer.' + optimizer).__getattribute__('Optimizer')
self.__optimizer__ = Optimizer(self.__model__.parameters(), **kwargs)
Scheduler = importlib.import_module(
'scheduler.' + scheduler).__getattribute__('Scheduler')
self.__scheduler__, self.lr_step = Scheduler(self.__optimizer__,
**kwargs)
self.scaler = GradScaler()
self.gpu = gpu
self.mixedprec = mixedprec
assert self.lr_step in ['epoch', 'iteration']
# ## ===== ===== ===== ===== ===== ===== ===== =====
# ## Train network
# ## ===== ===== ===== ===== ===== ===== ===== =====
def train_network(self, loader, verbose):
self.__model__.train()
stepsize = loader.batch_size
counter = 0
index = 0
loss = 0
top1 = 0 # EER or accuracy
tstart = time.time()
for data, data_label in loader:
data = data.transpose(1, 0)
self.__model__.zero_grad()
label = torch.LongTensor(data_label).cuda()
if self.mixedprec:
with autocast():
nloss, prec1 = self.__model__(data, label)
self.scaler.scale(nloss).backward()
self.scaler.step(self.__optimizer__)
self.scaler.update()
else:
nloss, prec1 = self.__model__(data, label)
nloss.backward()
self.__optimizer__.step()
loss += nloss.detach().cpu()
top1 += prec1
counter += 1
index += stepsize
telapsed = time.time() - tstart
tstart = time.time()
if verbose:
sys.stdout.write("\rProcessing (%d) " % (index))
sys.stdout.write(
"Loss %f TEER/TAcc %2.3f%% - %.2f Hz " %
(loss / counter, top1 / counter, stepsize / telapsed))
sys.stdout.flush()
if self.lr_step == 'iteration': self.__scheduler__.step()
if self.lr_step == 'epoch': self.__scheduler__.step()
sys.stdout.write("\n")
return (loss / counter, top1 / counter)
## ===== ===== ===== ===== ===== ===== ===== =====
## Evaluate from list
## ===== ===== ===== ===== ===== ===== ===== =====
def evaluateFromList(self,
test_list,
test_path,
nDataLoaderThread,
print_interval=100,
num_eval=10,
**kwargs):
self.__model__.eval()
lines = []
files = []
feats = {}
tstart = time.time()
## Read all lines
with open(test_list) as f:
lines = f.readlines()
## Get a list of unique file names
files = sum([x.strip().split()[-2:] for x in lines], [])
setfiles = list(set(files))
setfiles.sort()
## Define test data loader
test_dataset = test_dataset_loader(setfiles,
test_path,
num_eval=num_eval,
**kwargs)
test_loader = torch.utils.data.DataLoader(
test_dataset,
batch_size=1,
shuffle=False,
num_workers=nDataLoaderThread,
drop_last=False,
)
## Extract features for every image
for idx, data in enumerate(test_loader):
inp1 = data[0][0].cuda()
ref_feat = self.__model__(inp1).detach().cpu()
feats[data[1][0]] = ref_feat
telapsed = time.time() - tstart
if idx % print_interval == 0:
sys.stdout.write(
"\rReading %d of %d: %.2f Hz, embedding size %d" %
(idx, len(setfiles), idx / telapsed, ref_feat.size()[1]))
print('')
all_scores = []
all_labels = []
all_trials = []
tstart = time.time()
## Read files and compute all scores
for idx, line in enumerate(lines):
data = line.split()
## Append random label if missing
if len(data) == 2: data = [random.randint(0, 1)] + data
ref_feat = feats[data[1]].cuda()
com_feat = feats[data[2]].cuda()
if self.__model__.module.__L__.test_normalize:
ref_feat = F.normalize(ref_feat, p=2, dim=1)
com_feat = F.normalize(com_feat, p=2, dim=1)
dist = F.pairwise_distance(ref_feat.unsqueeze(-1),
com_feat.unsqueeze(-1).transpose(
0, 2)).detach().cpu().numpy()
score = -1 * numpy.mean(dist)
all_scores.append(score)
all_labels.append(int(data[0]))
all_trials.append(data[1] + " " + data[2])
if idx % print_interval == 0:
telapsed = time.time() - tstart
sys.stdout.write("\rComputing %d of %d: %.2f Hz" %
(idx, len(lines), idx / telapsed))
sys.stdout.flush()
print('')
return (all_scores, all_labels, all_trials)
## ===== ===== ===== ===== ===== ===== ===== =====
## Save parameters
## ===== ===== ===== ===== ===== ===== ===== =====
def saveParameters(self, path):
torch.save(self.__model__.module.state_dict(), path)
## ===== ===== ===== ===== ===== ===== ===== =====
## Load parameters
## ===== ===== ===== ===== ===== ===== ===== =====
def loadParameters(self, path):
self_state = self.__model__.module.state_dict()
loaded_state = torch.load(path, map_location="cuda:%d" % self.gpu)
for name, param in loaded_state.items():
origname = name
if name not in self_state:
name = name.replace("module.", "")
if name not in self_state:
print("%s is not in the model." % origname)
continue
if self_state[name].size() != loaded_state[origname].size():
print("Wrong parameter length: %s, model: %s, loaded: %s" %
(origname, self_state[name].size(),
loaded_state[origname].size()))
continue
self_state[name].copy_(param)
def train(fold):
args = get_args()
with open(args.config) as file:
config_file = yaml.load(file, Loader=yaml.FullLoader)
wandb.init(
project="siim2020",
entity="siim_melanoma",
# name=f"20200718-effb0-adamw-consineaneal-{fold}",
name=f"2017-2018-rexnet-test-{fold}",
#name=f"swav-test-{fold}",
#name=f"RAdam-b6-384x384-{fold}"
)
config = wandb.config # Initialize config
config.update(config_file)
device = config.device
model_path = config.model_path.format(fold)
seed_everything(config.seed)
df = pd.read_csv(config.train_csv_fold)
df_train = df[df.kfold != fold].reset_index(drop=True)
df_train["image_name"] = config.training_data_path + df_train[
"image_name"] + ".jpg"
if config.supplement_data["use_supplement"]:
print(f"training shape before merge {df_train.shape}")
df_supplement = pd.read_csv(config.supplement_data["csv_file"])
df_supplement = df_supplement[df_supplement["tfrecord"] % 2 == 0]
df_supplement = df_supplement[df_supplement["target"] == 1]
df_supplement["image_name"] = (config.supplement_data["file_path"] +
df_supplement["image_name"] + ".jpg")
df_train = pd.concat([df_train, df_supplement]).reset_index(drop=True)
df_train = df_train.sample(
frac=1, random_state=config.seed).reset_index(drop=True)
del df_supplement
print(f"training shape after merge {df_train.shape}")
df_valid = df[df.kfold == fold].reset_index(drop=True)
df_valid["image_name"] = config.training_data_path + df_valid[
"image_name"] + ".jpg"
if config.use_metadata:
df_train, meta_features = get_meta_feature(df_train)
df_valid, _ = get_meta_feature(df_valid)
else:
meta_features = None
model = get_model(
config.model_backbone,
config.model_name,
config.num_classes,
config.input_size,
config.use_metadata,
meta_features,
)
model = model.to(config.device)
print("watching model")
wandb.watch(model, log="all")
mean = (0.485, 0.456, 0.406)
std = (0.229, 0.224, 0.225)
train_aug = albumentations.Compose([
AdvancedHairAugmentation(hairs_folder="../input/melanoma-hairs/"),
# albumentations.augmentations.transforms.CenterCrop(64, 64, p=0.8),
albumentations.augmentations.transforms.RandomBrightnessContrast(),
albumentations.augmentations.transforms.HueSaturationValue(),
# Microscope(p=0.4),
albumentations.augmentations.transforms.RandomResizedCrop(
config.input_size, config.input_size, scale=(0.7, 1.0), p=0.4),
albumentations.augmentations.transforms.VerticalFlip(p=0.4),
albumentations.augmentations.transforms.Cutout(p=0.3), # doesnt work
albumentations.ShiftScaleRotate(shift_limit=0.0625,
scale_limit=0.1,
rotate_limit=15),
albumentations.Flip(p=0.5),
RandomAugMix(severity=7, width=7, alpha=5, p=0.3),
# albumentations.augmentations.transforms.Resize(
# config.input_size, config.input_size, p=1
# ),
albumentations.Normalize(mean,
std,
max_pixel_value=255.0,
always_apply=True),
])
valid_aug = albumentations.Compose([
albumentations.Normalize(mean,
std,
max_pixel_value=255.0,
always_apply=True),
])
train_images = df_train.image_name.values.tolist()
# train_images = [
# os.path.join(config.training_data_path, i + ".jpg") for i in train_images
# ]
train_targets = df_train.target.values
valid_images = df_valid.image_name.values.tolist()
# valid_images = [
# os.path.join(config.training_data_path, i + ".jpg") for i in valid_images
# ]
valid_targets = df_valid.target.values
train_dataset = ClassificationDataset(
image_paths=train_images,
targets=train_targets,
resize=None,
augmentations=train_aug,
meta_features=meta_features,
df_meta_features=df_train,
)
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=config.batch_size,
# num_workers=4,
num_workers=1,
pin_memory=True,
shuffle=True,
#sampler=BalanceClassSampler(labels=train_targets, mode="upsampling"),
drop_last=True,
)
valid_dataset = ClassificationDataset(
image_paths=valid_images,
targets=valid_targets,
resize=None,
augmentations=valid_aug,
meta_features=meta_features,
df_meta_features=df_valid,
)
valid_loader = torch.utils.data.DataLoader(
valid_dataset,
batch_size=config.test_batch_size,
shuffle=False,
# num_workers=4,
num_workers=1,
pin_memory=True,
# drop_last=True
)
#optimizer = torch.optim.AdamW(model.parameters(), lr=config.lr)
optimizer = RAdam(model.parameters(), lr=config.lr)
if config.swa["use_swa"]:
optimizer = SWA(optimizer, swa_start=12, swa_freq=1)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer,
patience=2,
threshold=0.0001,
mode="max")
# scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
# optimizer, len(train_loader) * config.epochs
# )
#scheduler = torch.optim.lr_scheduler.CyclicLR(
# optimizer,
# base_lr=config.lr / 10,
# max_lr=config.lr * 100,
# mode="triangular2",
# cycle_momentum=False,
#)
#scheduler = torch.optim.lr_scheduler.OneCycleLR(
# optimizer, max_lr=3e-3, steps_per_epoch=len(train_loader), epochs=config.epochs
#)
es = EarlyStopping(patience=6, mode="max")
if config.fp16:
print("************* using fp16 *************")
scaler = GradScaler()
else:
scaler = False
for epoch in range(config.epochs):
train_loss = Engine.train(
train_loader,
model,
optimizer,
device=config.device,
wandb=wandb,
accumulation_steps=config.accumulation_steps,
fp16=config.fp16,
scaler=scaler,
)
predictions, valid_loss = Engine.evaluate(
valid_loader,
model,
device=config.device,
wandb=wandb,
epoch=epoch,
upload_image=False,
use_sigmoid=True,
)
predictions = np.vstack((predictions)).ravel()
auc = metrics.roc_auc_score(valid_targets, predictions)
print(f"Epoch = {epoch}, AUC = {auc}")
wandb.log({
"valid_auc": auc,
})
scheduler.step(auc)
es(auc, model, model_path=model_path)
if es.early_stop:
print("Early stopping")
break
if config.swa["use_swa"]:
print("saving the model using SWA")
optimizer.swap_swa_sgd()
torch.save(model.state_dict(), config.swa["model_path"].format(fold))
evaluate_for_best_epoch(
fold,
model_path,
config.device,
valid_loader,
config.model_name,
valid_targets,
"final",
meta_features=meta_features,
)
if config.swa["use_swa"]:
model_path = config.swa["model_path"].format(fold)
evaluate_for_best_epoch(
fold,
model_path,
config.device,
valid_loader,
config.model_name,
valid_targets,
"swa",
meta_features=meta_features,
)
def train_job(model_name, train_df, valid_df, model_ckpt=None, log=True):
if log:
neptune.set_project("utsav/wheat-det")
neptune.init("utsav/wheat-det", api_token=NEPTUNE_API_TOKEN)
neptune.create_experiment(
FLAGS["exp_name"],
exp_description,
params=FLAGS,
upload_source_files="*.txt",
)
best_score = 0.0
start_epoch = 0
datasets = get_training_datasets(train_df, valid_df)
train_loader = DataLoader(
datasets["train"],
batch_size=FLAGS["batch_size"],
num_workers=FLAGS["num_workers"],
shuffle=True, # sampler=sampler, #
collate_fn=collate_fn,
)
val_loader = DataLoader(
datasets["valid"],
batch_size=FLAGS["batch_size"] * 2,
shuffle=False,
num_workers=FLAGS["num_workers"],
collate_fn=collate_fn,
)
if model_ckpt is not None:
model = get_train_model(model_name, model_ckpt)
else:
model = get_train_model(model_name)
model.to(device)
optimizer = Ranger(
model.parameters(),
lr=FLAGS["learning_rate"],
alpha=0.5,
k=6,
N_sma_threshhold=5,
weight_decay=FLAGS["weight_decay"],
)
scheduler = lr_scheduler.ReduceLROnPlateau(
optimizer,
"min",
factor=0.5,
verbose=True,
patience=FLAGS["scheduler_pat"],
)
scaler = GradScaler()
es = 0
for epoch in range(start_epoch, FLAGS["num_epochs"]):
print("-" * 27 + f"Epoch #{epoch+1} started" + "-" * 27)
train_loss = train_one_epoch(
train_loader,
model,
optimizer,
device,
scaler,
scheduler=None,
log=log,
)
print(f"Average loss for epoch #{epoch+1} : {train_loss}")
val_metric, val_loss = val_one_epoch(model, val_loader)
scheduler.step(val_loss)
print(f"metric/val : {val_metric}")
print(f"loss/val : {val_loss}")
if log:
neptune.log_metric("metric/val", val_metric)
neptune.log_metric("loss/val", val_loss)
# if epoch==FLAGS['unfreeze_epoch']:
# model = unfreeze_all_layers(model)
if (val_metric > best_score) or (best_score - val_metric < 0.01):
es = 0
if val_metric > best_score:
best_score = val_metric
if epoch > 9:
save_upload(
model,
optimizer,
epoch,
scheduler,
val_metric,
exp_name=FLAGS["exp_name"],
)
# else:
# if epoch>24:
# es+=1
# if es > FLAGS['early_stop_count']:
# print('Early stopping...')
# break
print("-" * 28 + f"Epoch #{epoch+1} ended" + "-" * 28)
neptune.stop()
def train(self, args, logger=None):
"""
Train function of FixMatch.
From data_loader, it inference training data, computes losses, and update the networks.
"""
ngpus_per_node = torch.cuda.device_count()
# lb: labeled, ulb: unlabeled
self.train_model.train()
# for gpu profiling
start_batch = torch.cuda.Event(enable_timing=True)
end_batch = torch.cuda.Event(enable_timing=True)
start_run = torch.cuda.Event(enable_timing=True)
end_run = torch.cuda.Event(enable_timing=True)
start_batch.record()
best_eval_acc, best_it = 0.0, 0
scaler = GradScaler()
amp_cm = autocast if args.amp else contextlib.nullcontext
for (x_lb, y_lb, idx), (x_ulb_w, x_ulb_s,
_) in zip(self.loader_dict['train_lb'],
self.loader_dict['train_ulb']):
# prevent the training iterations exceed args.num_train_iter
if self.it > args.num_train_iter:
break
end_batch.record()
torch.cuda.synchronize()
start_run.record()
num_lb = x_lb.shape[0]
num_ulb = x_ulb_w.shape[0]
assert num_ulb == x_ulb_s.shape[0]
x_lb, x_ulb_w, x_ulb_s = x_lb.cuda(args.gpu), x_ulb_w.cuda(
args.gpu), x_ulb_s.cuda(args.gpu)
y_lb = y_lb.cuda(args.gpu)
idx = idx.cuda(args.gpu)
inputs = torch.cat((x_lb, x_ulb_w, x_ulb_s))
# inference and calculate sup/unsup losses
with amp_cm():
# logits = self.train_model(inputs)
# logits_x_lb = logits[:num_lb]
# logits_x_ulb_w, logits_x_ulb_s = logits[num_lb:].chunk(2)
# del logits
logits_x_lb, logits_x_ulb_w, logits_x_ulb_s, sup_soft_labels = self.train_model(
inputs, num_lb=num_lb, it=self.it, sup_index=idx)
# hyper-params for update
T = self.t_fn(self.it)
p_cutoff = self.p_fn(self.it)
if sup_soft_labels is not None:
sup_loss = sat_loss(logits_x_lb, sup_soft_labels)
else:
sup_loss = ce_loss(logits_x_lb, y_lb, reduction='mean')
unsup_loss, mask = consistency_loss(
logits_x_ulb_w,
logits_x_ulb_s,
'ce',
T,
p_cutoff,
use_hard_labels=args.hard_label)
total_loss = sup_loss + self.lambda_u * unsup_loss
# parameter updates
if args.amp:
scaler.scale(total_loss).backward()
scaler.step(self.optimizer)
scaler.update()
else:
total_loss.backward()
self.optimizer.step()
self.scheduler.step()
self.train_model.zero_grad()
with torch.no_grad():
self._eval_model_update()
end_run.record()
torch.cuda.synchronize()
# tensorboard_dict update
tb_dict = {}
tb_dict['train/sup_loss'] = sup_loss.detach()
tb_dict['train/unsup_loss'] = unsup_loss.detach()
tb_dict['train/total_loss'] = total_loss.detach()
tb_dict['train/mask_ratio'] = 1.0 - mask.detach()
tb_dict['lr'] = self.optimizer.param_groups[0]['lr']
tb_dict['train/prefecth_time'] = start_batch.elapsed_time(
end_batch) / 1000.
tb_dict['train/run_time'] = start_run.elapsed_time(end_run) / 1000.
if self.it % self.num_eval_iter == 0:
eval_dict = self.evaluate(args=args)
tb_dict.update(eval_dict)
save_path = os.path.join(args.save_dir, args.save_name)
if tb_dict['eval/top-1-acc'] > best_eval_acc:
best_eval_acc = tb_dict['eval/top-1-acc']
best_it = self.it
self.print_fn(
f"{self.it} iteration, USE_EMA: {hasattr(self, 'eval_model')}, {tb_dict}, BEST_EVAL_ACC: {best_eval_acc}, at {best_it} iters"
)
if not args.multiprocessing_distributed or \
(args.multiprocessing_distributed and args.rank % ngpus_per_node == 0):
if self.it % self.num_eval_iter == 0:
self.save_model('model_last.pth', save_path)
if self.it % self.num_eval_iter == 0:
self.save_model('model_{}.pth'.format(self.it), save_path)
if self.it == best_it:
self.save_model('model_best.pth', save_path)
if not self.tb_log is None:
self.tb_log.update(tb_dict, self.it)
self.it += 1
del tb_dict
start_batch.record()
if self.it > 2**19:
self.num_eval_iter = 1000
eval_dict = self.evaluate(args=args)
eval_dict.update({
'eval/best_acc': best_eval_acc,
'eval/best_it': best_it
})
return eval_dict
def _train(self, rank):
_seed_everything(self._data_shuffle_seed)
self._setup_ddp(rank)
self._rank = rank
self._scaler = GradScaler()
self._model = self._get_model(self._rank)
self._optimizer = AdamW(params=self._model.parameters(),
lr=self._learning_rate)
self._global_step = 0
self._samples_seen = 0
self._train_dl = self._get_dataloader(is_train=True, samples_offset=0)
self._valid_dl = self._get_dataloader(is_train=False, samples_offset=0)
steps_per_epoch = len(self._train_dl)
num_training_steps = steps_per_epoch * self._n_epochs
num_warmup_steps = self._warmup_ratio * num_training_steps
self._scheduler = get_linear_schedule_with_warmup(
optimizer=self._optimizer,
num_warmup_steps=num_warmup_steps,
num_training_steps=num_training_steps)
if self._checkpoint_file_path.is_file():
self._load_checkpoint()
elif self._init_weights_from_checkpoint:
self._load_only_weights_from_checkpoint()
while True:
if self._rank == 0:
self._writer = self._writer or SummaryWriter(
self._experiment_dir / 'tb_logs')
self._train_dl = tqdm.tqdm(self._train_dl,
desc='Train step',
total=num_training_steps,
position=1,
initial=self._global_step)
for i_step, model_input in enumerate(self._train_dl):
train_losses_dict = self._train_step(model_input)
if rank == 0:
self._train_dl.set_postfix({
'samples_seen':
self._samples_seen,
'epoch':
self._global_step / steps_per_epoch
})
self._write_tb_logs(train_losses_dict)
self._write_tb_logs({
'learning-rate':
self._optimizer.param_groups[0]['lr']
})
self._write_tb_logs(
{'max_seq_len': model_input.input_ids.size()[1]})
if self._rank == 0 and self._global_step % self._validate_each_n_steps == 0:
valid_loss = self._validate()
self._save_checkpoint()
self._write_tb_logs({'loss/valid': valid_loss})
if self._global_step >= num_training_steps:
break
dist.destroy_process_group()
class Trainer:
"""
Implements the training logic. Some common configuration (checkpointing frequency, path, validation frequency)
is done by checking util.common_opts that is set via the CL.
"""
def __init__(
self,
game: torch.nn.Module,
optimizer: torch.optim.Optimizer,
train_data: DataLoader,
optimizer_scheduler: Optional[torch.optim.lr_scheduler._LRScheduler] = None,
validation_data: Optional[DataLoader] = None,
device: torch.device = None,
callbacks: Optional[List[Callback]] = None,
grad_norm: float = None,
aggregate_interaction_logs: bool = True,
):
"""
:param game: A nn.Module that implements forward(); it is expected that forward returns a tuple of (loss, d),
where loss is differentiable loss to be minimized and d is a dictionary (potentially empty) with auxiliary
metrics that would be aggregated and reported
:param optimizer: An instance of torch.optim.Optimizer
:param optimizer_scheduler: An optimizer scheduler to adjust lr throughout training
:param train_data: A DataLoader for the training set
:param validation_data: A DataLoader for the validation set (can be None)
:param device: A torch.device on which to tensors should be stored
:param callbacks: A list of egg.core.Callback objects that can encapsulate monitoring or checkpointing
"""
self.game = game
self.optimizer = optimizer
self.optimizer_scheduler = optimizer_scheduler
self.train_data = train_data
self.validation_data = validation_data
common_opts = get_opts()
self.validation_freq = common_opts.validation_freq
self.device = common_opts.device if device is None else device
self.should_stop = False
self.start_epoch = 0 # Can be overwritten by checkpoint loader
self.callbacks = callbacks if callbacks else []
self.grad_norm = grad_norm
self.aggregate_interaction_logs = aggregate_interaction_logs
self.update_freq = common_opts.update_freq
if common_opts.load_from_checkpoint is not None:
print(
f"# Initializing model, trainer, and optimizer from {common_opts.load_from_checkpoint}"
)
self.load_from_checkpoint(common_opts.load_from_checkpoint)
self.distributed_context = common_opts.distributed_context
if self.distributed_context.is_distributed:
print("# Distributed context: ", self.distributed_context)
if self.distributed_context.is_leader and not any(
isinstance(x, CheckpointSaver) for x in self.callbacks
):
if common_opts.preemptable:
assert (
common_opts.checkpoint_dir
), "checkpointing directory has to be specified"
d = get_preemptive_checkpoint_dir(common_opts.checkpoint_dir)
self.checkpoint_path = d
self.load_from_latest(d)
else:
self.checkpoint_path = (
None
if common_opts.checkpoint_dir is None
else pathlib.Path(common_opts.checkpoint_dir)
)
if self.checkpoint_path:
checkpointer = CheckpointSaver(
checkpoint_path=self.checkpoint_path,
checkpoint_freq=common_opts.checkpoint_freq,
)
self.callbacks.append(checkpointer)
if self.distributed_context.is_leader and common_opts.tensorboard:
assert (
common_opts.tensorboard_dir
), "tensorboard directory has to be specified"
tensorboard_logger = TensorboardLogger()
self.callbacks.append(tensorboard_logger)
if self.callbacks is None:
self.callbacks = [
ConsoleLogger(print_train_loss=False, as_json=False),
]
if self.distributed_context.is_distributed:
device_id = self.distributed_context.local_rank
torch.cuda.set_device(device_id)
self.game.to(device_id)
# NB: here we are doing something that is a bit shady:
# 1/ optimizer was created outside of the Trainer instance, so we don't really know
# what parameters it optimizes. If it holds something what is not within the Game instance
# then it will not participate in distributed training
# 2/ if optimizer only holds a subset of Game parameters, it works, but somewhat non-documentedly.
# In fact, optimizer would hold parameters of non-DistributedDataParallel version of the Game. The
# forward/backward calls, however, would happen on the DistributedDataParallel wrapper.
# This wrapper would sync gradients of the underlying tensors - which are the ones that optimizer
# holds itself. As a result it seems to work, but only because DDP doesn't take any tensor ownership.
self.game = torch.nn.parallel.DistributedDataParallel(
self.game,
device_ids=[device_id],
output_device=device_id,
find_unused_parameters=True,
)
self.optimizer.state = move_to(self.optimizer.state, device_id)
else:
self.game.to(self.device)
# NB: some optimizers pre-allocate buffers before actually doing any steps
# since model is placed on GPU within Trainer, this leads to having optimizer's state and model parameters
# on different devices. Here, we protect from that by moving optimizer's internal state to the proper device
self.optimizer.state = move_to(self.optimizer.state, self.device)
if common_opts.fp16:
self.scaler = GradScaler()
else:
self.scaler = None
def eval(self, data=None):
mean_loss = 0.0
interactions = []
n_batches = 0
validation_data = self.validation_data if data is None else data
self.game.eval()
with torch.no_grad():
for batch in validation_data:
if not isinstance(batch, Batch):
batch = Batch(*batch)
batch = batch.to(self.device)
optimized_loss, interaction = self.game(*batch)
if (
self.distributed_context.is_distributed
and self.aggregate_interaction_logs
):
interaction = Interaction.gather_distributed_interactions(
interaction
)
interaction = interaction.to("cpu")
mean_loss += optimized_loss
for callback in self.callbacks:
callback.on_batch_end(
interaction, optimized_loss, n_batches, is_training=False
)
interactions.append(interaction)
n_batches += 1
mean_loss /= n_batches
full_interaction = Interaction.from_iterable(interactions)
return mean_loss.item(), full_interaction
def train_epoch(self):
mean_loss = 0
n_batches = 0
interactions = []
self.game.train()
self.optimizer.zero_grad()
for batch_id, batch in enumerate(self.train_data):
if not isinstance(batch, Batch):
batch = Batch(*batch)
batch = batch.to(self.device)
context = autocast() if self.scaler else nullcontext()
with context:
optimized_loss, interaction = self.game(*batch)
if self.update_freq > 1:
# throughout EGG, we minimize _mean_ loss, not sum
# hence, we need to account for that when aggregating grads
optimized_loss = optimized_loss / self.update_freq
if self.scaler:
self.scaler.scale(optimized_loss).backward()
else:
optimized_loss.backward()
if batch_id % self.update_freq == self.update_freq - 1:
if self.scaler:
self.scaler.unscale_(self.optimizer)
if self.grad_norm:
torch.nn.utils.clip_grad_norm_(
self.game.parameters(), self.grad_norm
)
if self.scaler:
self.scaler.step(self.optimizer)
self.scaler.update()
else:
self.optimizer.step()
self.optimizer.zero_grad()
n_batches += 1
mean_loss += optimized_loss.detach()
if (
self.distributed_context.is_distributed
and self.aggregate_interaction_logs
):
interaction = Interaction.gather_distributed_interactions(interaction)
interaction = interaction.to("cpu")
for callback in self.callbacks:
callback.on_batch_end(interaction, optimized_loss, batch_id)
interactions.append(interaction)
if self.optimizer_scheduler:
self.optimizer_scheduler.step()
mean_loss /= n_batches
full_interaction = Interaction.from_iterable(interactions)
return mean_loss.item(), full_interaction
def train(self, n_epochs):
for callback in self.callbacks:
callback.on_train_begin(self)
for epoch in range(self.start_epoch, n_epochs):
for callback in self.callbacks:
callback.on_epoch_begin(epoch + 1)
train_loss, train_interaction = self.train_epoch()
for callback in self.callbacks:
callback.on_epoch_end(train_loss, train_interaction, epoch + 1)
validation_loss = validation_interaction = None
if (
self.validation_data is not None
and self.validation_freq > 0
and (epoch + 1) % self.validation_freq == 0
):
for callback in self.callbacks:
callback.on_validation_begin(epoch + 1)
validation_loss, validation_interaction = self.eval()
for callback in self.callbacks:
callback.on_validation_end(
validation_loss, validation_interaction, epoch + 1
)
if self.should_stop:
for callback in self.callbacks:
callback.on_early_stopping(
train_loss,
train_interaction,
epoch + 1,
validation_loss,
validation_interaction,
)
break
for callback in self.callbacks:
callback.on_train_end()
def load(self, checkpoint: Checkpoint):
self.game.load_state_dict(checkpoint.model_state_dict)
self.optimizer.load_state_dict(checkpoint.optimizer_state_dict)
if checkpoint.optimizer_scheduler_state_dict:
self.optimizer_scheduler.load_state_dict(
checkpoint.optimizer_scheduler_state_dict
)
self.start_epoch = checkpoint.epoch
def load_from_checkpoint(self, path):
"""
Loads the game, agents, and optimizer state from a file
:param path: Path to the file
"""
print(f"# loading trainer state from {path}")
checkpoint = torch.load(path)
self.load(checkpoint)
def load_from_latest(self, path):
latest_file, latest_time = None, None
for file in path.glob("*.tar"):
creation_time = os.stat(file).st_ctime
if latest_time is None or creation_time > latest_time:
latest_file, latest_time = file, creation_time
if latest_file is not None:
self.load_from_checkpoint(latest_file)
class ConfigurableStep(Module):
def __init__(self, opt_step, env):
super(ConfigurableStep, self).__init__()
self.step_opt = opt_step
self.env = env
self.opt = env['opt']
self.gen_outputs = opt_step['generator_outputs']
self.loss_accumulator = LossAccumulator()
self.optimizers = None
self.scaler = GradScaler(enabled=self.opt['fp16'])
self.grads_generated = False
self.min_total_loss = opt_step['min_total_loss'] if 'min_total_loss' in opt_step.keys() else -999999999
self.clip_grad_eps = opt_get(opt_step, ['clip_grad_eps'], None)
# This is a half-measure that can be used between anomaly_detection and running a potentially problematic
# trainer bare. With this turned on, the optimizer will not step() if a nan grad is detected. If a model trips
# this warning 10 times in a row, the training session is aborted and the model state is saved. This has a
# noticeable affect on training speed, but nowhere near as bad as anomaly_detection.
self.check_grads_for_nan = opt_get(opt_step, ['check_grads_for_nan'], False)
self.nan_counter = 0
self.injectors = []
if 'injectors' in self.step_opt.keys():
injector_names = []
for inj_name, injector in self.step_opt['injectors'].items():
assert inj_name not in injector_names # Repeated names are always an error case.
injector_names.append(inj_name)
self.injectors.append(create_injector(injector, env))
losses = []
self.weights = {}
if 'losses' in self.step_opt.keys():
for loss_name, loss in self.step_opt['losses'].items():
assert loss_name not in self.weights.keys() # Repeated names are always an error case.
losses.append((loss_name, create_loss(loss, env)))
self.weights[loss_name] = loss['weight']
self.losses = OrderedDict(losses)
def get_network_for_name(self, name):
return self.env['generators'][name] if name in self.env['generators'].keys() \
else self.env['discriminators'][name]
# Subclasses should override this to define individual optimizers. They should all go into self.optimizers.
# This default implementation defines a single optimizer for all Generator parameters.
# Must be called after networks are initialized and wrapped.
def define_optimizers(self):
opt_configs = [opt_get(self.step_opt, ['optimizer_params'], None)]
self.optimizers = []
if opt_configs[0] is None:
return
training = self.step_opt['training']
training_net = self.get_network_for_name(training)
nets = [training_net]
training = [training]
for net_name, net, opt_config in zip(training, nets, opt_configs):
# Configs can organize parameters by-group and specify different learning rates for each group. This only
# works in the model specifically annotates which parameters belong in which group using PARAM_GROUP.
optim_params = {'default': {'params': [], 'lr': opt_config['lr']}}
if opt_config is not None and 'param_groups' in opt_config.keys():
for k, pg in opt_config['param_groups'].items():
optim_params[k] = {'params': [], 'lr': pg['lr']}
for k, v in net.named_parameters(): # can optimize for a part of the model
# Make some inference about these parameters, which can be used by some optimizers to treat certain
# parameters differently. For example, it is considered good practice to not do weight decay on
# BN & bias parameters. TODO: process the module tree instead of the parameter tree to accomplish the
# same thing, but in a more effective way.
if k.endswith(".bias"):
v.is_bias = True
if k.endswith(".weight"):
v.is_weight = True
if ".bn" in k or '.batchnorm' in k or '.bnorm' in k:
v.is_bn = True
# Some models can specify some parameters to be in different groups.
param_group = "default"
if hasattr(v, 'PARAM_GROUP'):
if v.PARAM_GROUP in optim_params.keys():
param_group = v.PARAM_GROUP
else:
logger.warning(f'Model specifies a custom param group {v.PARAM_GROUP} which is not configured. '
f'The same LR will be used for all parameters.')
if v.requires_grad:
optim_params[param_group]['params'].append(v)
else:
if self.env['rank'] <= 0:
logger.warning('Params [{:s}] will not optimize.'.format(k))
if 'optimizer' not in self.step_opt.keys() or self.step_opt['optimizer'] == 'adam':
opt = torch.optim.Adam(list(optim_params.values()),
weight_decay=opt_config['weight_decay'],
betas=(opt_config['beta1'], opt_config['beta2']))
elif self.step_opt['optimizer'] == 'adamw':
opt = torch.optim.AdamW(list(optim_params.values()),
weight_decay=opt_get(opt_config, ['weight_decay'], 1e-2),
betas=(opt_get(opt_config, ['beta1'], .9), opt_get(opt_config, ['beta2'], .999)))
elif self.step_opt['optimizer'] == 'lars':
from trainer.optimizers.larc import LARC
from trainer.optimizers.sgd import SGDNoBiasMomentum
optSGD = SGDNoBiasMomentum(list(optim_params.values()), lr=opt_config['lr'], momentum=opt_config['momentum'],
weight_decay=opt_config['weight_decay'])
opt = LARC(optSGD, trust_coefficient=opt_config['lars_coefficient'])
elif self.step_opt['optimizer'] == 'sgd':
from torch.optim import SGD
opt = SGD(list(optim_params.values()), lr=opt_config['lr'], momentum=opt_config['momentum'], weight_decay=opt_config['weight_decay'])
opt._config = opt_config # This is a bit seedy, but we will need these configs later.
opt._config['network'] = net_name
self.optimizers.append(opt)
# Returns all optimizers used in this step.
def get_optimizers(self):
assert self.optimizers is not None
return self.optimizers
# Returns optimizers which are opting in for default LR scheduling.
def get_optimizers_with_default_scheduler(self):
assert self.optimizers is not None
return self.optimizers
# Returns the names of the networks this step will train. Other networks will be frozen.
def get_networks_trained(self):
if isinstance(self.step_opt['training'], list):
return self.step_opt['training']
else:
return [self.step_opt['training']]
def get_training_network_name(self):
if isinstance(self.step_opt['training'], list):
return self.step_opt['training'][0]
else:
return self.step_opt['training']
# Performs all forward and backward passes for this step given an input state. All input states are lists of
# chunked tensors. Use grad_accum_step to dereference these steps. Should return a dict of tensors that later
# steps might use. These tensors are automatically detached and accumulated into chunks.
def do_forward_backward(self, state, grad_accum_step, amp_loss_id, train=True, no_ddp_sync=False):
local_state = {} # <-- Will store the entire local state to be passed to injectors & losses.
new_state = {} # <-- Will store state values created by this step for returning to ExtensibleTrainer.
for k, v in state.items():
local_state[k] = v[grad_accum_step]
local_state['train_nets'] = str(self.get_networks_trained())
# Some losses compute backward() internally. Accommodate this by stashing the amp_loss_id in env.
self.env['amp_loss_id'] = amp_loss_id
self.env['current_step_optimizers'] = self.optimizers
self.env['training'] = train
# Inject in any extra dependencies.
for inj in self.injectors:
# Don't do injections tagged with eval unless we are not in train mode.
if train and 'eval' in inj.opt.keys() and inj.opt['eval']:
continue
# Likewise, don't do injections tagged with train unless we are not in eval.
if not train and 'train' in inj.opt.keys() and inj.opt['train']:
continue
# Don't do injections tagged with 'after' or 'before' when we are out of spec.
if 'after' in inj.opt.keys() and self.env['step'] < inj.opt['after'] or \
'before' in inj.opt.keys() and self.env['step'] > inj.opt['before'] or \
'every' in inj.opt.keys() and self.env['step'] % inj.opt['every'] != 0:
continue
if 'no_accum' in inj.opt.keys() and grad_accum_step > 0:
continue
training_net = self.get_network_for_name(self.step_opt['training'])
if no_ddp_sync and hasattr(training_net, 'no_sync'):
with training_net.no_sync():
injected = inj(local_state)
else:
injected = inj(local_state)
local_state.update(injected)
new_state.update(injected)
if train and len(self.losses) > 0:
# Finally, compute the losses.
total_loss = 0
for loss_name, loss in self.losses.items():
# Some losses only activate after a set number of steps. For example, proto-discriminator losses can
# be very disruptive to a generator.
if 'after' in loss.opt.keys() and loss.opt['after'] > self.env['step'] or \
'before' in loss.opt.keys() and self.env['step'] > loss.opt['before'] or \
'every' in loss.opt.keys() and self.env['step'] % loss.opt['every'] != 0:
continue
if loss.is_stateful():
l, lstate = loss(self.get_network_for_name(self.step_opt['training']), local_state)
local_state.update(lstate)
new_state.update(lstate)
else:
l = loss(self.get_network_for_name(self.step_opt['training']), local_state)
total_loss += l * self.weights[loss_name]
# Record metrics.
if isinstance(l, torch.Tensor):
self.loss_accumulator.add_loss(loss_name, l)
for n, v in loss.extra_metrics():
self.loss_accumulator.add_loss("%s_%s" % (loss_name, n), v)
loss.clear_metrics()
# In some cases, the loss could not be set (e.g. all losses have 'after')
if isinstance(total_loss, torch.Tensor):
self.loss_accumulator.add_loss("%s_total" % (self.get_training_network_name(),), total_loss)
reset_required = total_loss < self.min_total_loss
# Scale the loss down by the accumulation factor.
total_loss = total_loss / self.env['mega_batch_factor']
# Get dem grads!
self.scaler.scale(total_loss).backward()
if reset_required:
# You might be scratching your head at this. Why would you zero grad as opposed to not doing a
# backwards? Because DDP uses the backward() pass as a synchronization point and there is not a good
# way to simply bypass backward. If you want a more efficient way to specify a min_loss, use or
# implement it at the loss level.
self.get_network_for_name(self.step_opt['training']).zero_grad()
self.loss_accumulator.increment_metric("%s_skipped_steps" % (self.get_training_network_name(),))
self.grads_generated = True
# Detach all state variables. Within the step, gradients can flow. Once these variables leave the step
# we must release the gradients.
new_state = recursively_detach(new_state)
return new_state
# Performs the optimizer step after all gradient accumulation is completed. Default implementation simply steps()
# all self.optimizers.
def do_step(self, step):
if not self.grads_generated:
return
self.grads_generated = False
for opt in self.optimizers:
# Optimizers can be opted out in the early stages of training.
after = opt._config['after'] if 'after' in opt._config.keys() else 0
after_network = self.opt['networks'][opt._config['network']]['after'] if 'after' in self.opt['networks'][opt._config['network']].keys() else 0
after = max(after, after_network)
if self.env['step'] < after:
continue
before = opt._config['before'] if 'before' in opt._config.keys() else -1
if before != -1 and self.env['step'] > before:
continue
nan_found = False
if self.check_grads_for_nan:
for pg in opt.param_groups:
for p in pg['params']:
if not torch.isfinite(p.grad).any():
nan_found = True
break
if nan_found:
break
if nan_found:
print("NaN found in grads. Throwing this step out.")
self.nan_counter += 1
else:
self.nan_counter = 0
if self.clip_grad_eps is not None:
for pg in opt.param_groups:
grad_norm = torch.nn.utils.clip_grad_norm_(pg['params'], self.clip_grad_eps)
if torch.isnan(grad_norm):
nan_found = True
self.nan_counter += 1
if not nan_found:
self.scaler.step(opt)
self.scaler.update()
def get_metrics(self):
return self.loss_accumulator.as_dict()
def __init__(
self,
game: torch.nn.Module,
optimizer: torch.optim.Optimizer,
train_data: DataLoader,
optimizer_scheduler: Optional[torch.optim.lr_scheduler._LRScheduler] = None,
validation_data: Optional[DataLoader] = None,
device: torch.device = None,
callbacks: Optional[List[Callback]] = None,
grad_norm: float = None,
aggregate_interaction_logs: bool = True,
):
"""
:param game: A nn.Module that implements forward(); it is expected that forward returns a tuple of (loss, d),
where loss is differentiable loss to be minimized and d is a dictionary (potentially empty) with auxiliary
metrics that would be aggregated and reported
:param optimizer: An instance of torch.optim.Optimizer
:param optimizer_scheduler: An optimizer scheduler to adjust lr throughout training
:param train_data: A DataLoader for the training set
:param validation_data: A DataLoader for the validation set (can be None)
:param device: A torch.device on which to tensors should be stored
:param callbacks: A list of egg.core.Callback objects that can encapsulate monitoring or checkpointing
"""
self.game = game
self.optimizer = optimizer
self.optimizer_scheduler = optimizer_scheduler
self.train_data = train_data
self.validation_data = validation_data
common_opts = get_opts()
self.validation_freq = common_opts.validation_freq
self.device = common_opts.device if device is None else device
self.should_stop = False
self.start_epoch = 0 # Can be overwritten by checkpoint loader
self.callbacks = callbacks if callbacks else []
self.grad_norm = grad_norm
self.aggregate_interaction_logs = aggregate_interaction_logs
self.update_freq = common_opts.update_freq
if common_opts.load_from_checkpoint is not None:
print(
f"# Initializing model, trainer, and optimizer from {common_opts.load_from_checkpoint}"
)
self.load_from_checkpoint(common_opts.load_from_checkpoint)
self.distributed_context = common_opts.distributed_context
if self.distributed_context.is_distributed:
print("# Distributed context: ", self.distributed_context)
if self.distributed_context.is_leader and not any(
isinstance(x, CheckpointSaver) for x in self.callbacks
):
if common_opts.preemptable:
assert (
common_opts.checkpoint_dir
), "checkpointing directory has to be specified"
d = get_preemptive_checkpoint_dir(common_opts.checkpoint_dir)
self.checkpoint_path = d
self.load_from_latest(d)
else:
self.checkpoint_path = (
None
if common_opts.checkpoint_dir is None
else pathlib.Path(common_opts.checkpoint_dir)
)
if self.checkpoint_path:
checkpointer = CheckpointSaver(
checkpoint_path=self.checkpoint_path,
checkpoint_freq=common_opts.checkpoint_freq,
)
self.callbacks.append(checkpointer)
if self.distributed_context.is_leader and common_opts.tensorboard:
assert (
common_opts.tensorboard_dir
), "tensorboard directory has to be specified"
tensorboard_logger = TensorboardLogger()
self.callbacks.append(tensorboard_logger)
if self.callbacks is None:
self.callbacks = [
ConsoleLogger(print_train_loss=False, as_json=False),
]
if self.distributed_context.is_distributed:
device_id = self.distributed_context.local_rank
torch.cuda.set_device(device_id)
self.game.to(device_id)
# NB: here we are doing something that is a bit shady:
# 1/ optimizer was created outside of the Trainer instance, so we don't really know
# what parameters it optimizes. If it holds something what is not within the Game instance
# then it will not participate in distributed training
# 2/ if optimizer only holds a subset of Game parameters, it works, but somewhat non-documentedly.
# In fact, optimizer would hold parameters of non-DistributedDataParallel version of the Game. The
# forward/backward calls, however, would happen on the DistributedDataParallel wrapper.
# This wrapper would sync gradients of the underlying tensors - which are the ones that optimizer
# holds itself. As a result it seems to work, but only because DDP doesn't take any tensor ownership.
self.game = torch.nn.parallel.DistributedDataParallel(
self.game,
device_ids=[device_id],
output_device=device_id,
find_unused_parameters=True,
)
self.optimizer.state = move_to(self.optimizer.state, device_id)
else:
self.game.to(self.device)
# NB: some optimizers pre-allocate buffers before actually doing any steps
# since model is placed on GPU within Trainer, this leads to having optimizer's state and model parameters
# on different devices. Here, we protect from that by moving optimizer's internal state to the proper device
self.optimizer.state = move_to(self.optimizer.state, self.device)
if common_opts.fp16:
self.scaler = GradScaler()
else:
self.scaler = None
def invert(
data_loader,
loss_fn,
optimizer,
steps=10,
scheduler=None,
use_amp=False,
grad_norm_fn=None,
callback_fn=None,
plot=False,
fig_path=None,
track_per_batch=False,
track_grad_norm=False,
):
assert valid_data_loader(
data_loader), f"invalid data_loader: {data_loader}"
params = sum((p_group['params'] for p_group in optimizer.param_groups), [])
lrs = [p_group['lr'] for p_group in optimizer.param_groups]
device = params[0].device
USE_AMP = (device.type == 'cuda') and use_amp
if USE_AMP:
scaler = GradScaler()
num_batches = len(data_loader)
track_len = steps * num_batches if track_per_batch else steps
metrics = pd.DataFrame({'step': [None] * track_len})
def process_result(res):
if isinstance(res, dict):
loss = res['loss']
info = res
for k, v in info.items():
info[k] = v.item() if isinstance(v, torch.Tensor) else v
elif isinstance(res, tuple):
loss, info = res
else:
loss = res
info = {'loss': loss.item()}
return loss, info
print(flush=True)
# if callback_fn:
# callback_fn(0, None)
with tqdmEpoch(steps, num_batches) as pbar:
for epoch in range(steps):
for batch_i, data in enumerate(data_loader):
optimizer.zero_grad()
if USE_AMP:
with autocast():
res = loss_fn(data)
loss, info = process_result(res)
scaler.scale(loss).backward()
grad_scale = scaler.get_scale()
else:
res = loss_fn(data)
loss, info = process_result(res)
loss.backward()
grad_scale = 1
if USE_AMP:
scaler.step(optimizer)
scaler.update()
else:
optimizer.step()
if scheduler is not None:
scheduler.step(loss)
if track_grad_norm or grad_norm_fn:
# XXX: probably shouldn't multiply with lr
total_norm = torch.norm(
torch.stack([
p.grad.detach().norm() / grad_scale # * lr
for p, lr in zip(params, lrs)
])).item()
if grad_norm_fn:
rescale_coef = grad_norm_fn(total_norm) / total_norm
for param in params:
param.grad.detach().mul_(rescale_coef)
info['|grad|'] = total_norm
pbar.set_postfix(**{
k: v
for k, v in info.items() if ']' not in k
},
refresh=False)
pbar.update()
if track_per_batch:
batch_total = epoch * num_batches + batch_i
step = batch_total
# step = epoch + (batch_i + 1) / num_batches
else:
step = epoch
# step = epoch + 1 + batch_i / num_batches
for k, v in info.items():
if k not in metrics: # add new column
metrics[k] = None
if metrics[k][step] is None:
metrics[k][step] = v
else:
metrics[k][step] += v
if not track_per_batch and batch_i == 0:
metrics['step'][epoch] = epoch + 1
if track_per_batch:
metrics['step'][batch_total] = (batch_total +
1) / num_batches
# batch end
if not track_per_batch:
for k, v in metrics.items():
if k != 'step':
metrics[k][epoch] /= num_batches
if callback_fn:
callback_fn(epoch + 1, metrics.iloc[step])
# epoch end
print(flush=True)
if plot and steps > 1:
plot_metrics(metrics, fig_path=fig_path, smoothing=0)
return metrics
def main(args):
torch.manual_seed(222)
torch.cuda.manual_seed_all(222)
np.random.seed(222)
print(args)
# config = []
device = torch.device('cuda')
model = None
if args.arch == "UNet":
model = UNet(args).to(device)
else:
raise ("architectures other than Unet hasn't been added!!")
# update_lrs = nn.Parameter(args.update_lr*torch.ones(self.update_step, len(self.net.vars)), requires_grad=True)
model.optimizer = optim.Adam(model.parameters(),
lr=args.lr,
eps=1e-7,
amsgrad=True,
weight_decay=args.weight_decay)
model.lr_scheduler = optim.lr_scheduler.ExponentialLR(
model.optimizer, args.exp_decay)
tmp = filter(lambda x: x.requires_grad, model.parameters())
num = sum(map(lambda x: np.prod(x.shape), tmp))
print(model)
#for name, param in model.named_parameters():
# print(name, param.size())
print('Total trainable tensors:', num, flush=True)
SUMMARY_INTERVAL = 5
TEST_PRINT_INTERVAL = SUMMARY_INTERVAL * 5
ITER_SAVE_INTERVAL = 300
EPOCH_SAVE_INTERVAL = 5
model_path = args.model_saving_path + args.model_name + "_batch_size_" + str(
args.batch_size) + "_lr_" + str(args.lr)
if not os.path.isdir(model_path):
os.mkdir(model_path)
ds = SimulationDataset(args.data_folder,
total_sample_number=args.total_sample_number)
means = [
1e-3, 1e-3
] #ds.means; #[Hy_meanR, Hy_meanI, Ex_meanR, Ex_meanI, Ez_meanR, Ez_meanI];
print("means: ", means)
torch.manual_seed(42)
train_ds, test_ds = random_split(
ds,
[int(0.9 * len(ds)), len(ds) - int(0.9 * len(ds))])
#print("total training samples: %d, total test samples: %d" % (len(train_ds), len(test_ds)), flush=True)
train_loader = DataLoader(train_ds,
batch_size=args.batch_size,
shuffle=True,
num_workers=0)
test_loader = DataLoader(test_ds,
batch_size=args.batch_size,
shuffle=True,
num_workers=0)
train_mean = 0
test_mean = 0
# first get the mean-absolute-field value:
for sample_batched in train_loader:
train_mean += torch.mean(torch.abs(sample_batched["field"]))
for sample_batched in test_loader:
test_mean += torch.mean(torch.abs(sample_batched["field"]))
train_mean /= len(train_loader)
test_mean /= len(test_loader)
print(
"total training samples: %d, total test samples: %d, train_abs_mean: %f, test_abs_mean: %f"
% (len(train_ds), len(test_ds), train_mean, test_mean),
flush=True)
# for visualizing the graph:
#writer = SummaryWriter('runs/'+args.model_name)
#test_input = None
#for sample in train_loader:
# test_input = sample['structure']
# break
#writer.add_graph(model, test_input.to(device))
#writer.close()
df = pd.DataFrame(columns=[
'epoch', 'train_loss', 'train_phys_reg', 'test_loss', 'test_phys_reg'
])
train_loss_history = []
train_phys_reg_history = []
test_loss_history = []
test_phys_reg_history = []
start_epoch = 0
if (args.continue_train):
print("Restoring weights from ",
model_path + "/last_model.pt",
flush=True)
checkpoint = torch.load(model_path + "/last_model.pt")
start_epoch = checkpoint['epoch']
model = checkpoint['model']
model.lr_scheduler = checkpoint['lr_scheduler']
model.optimizer = checkpoint['optimizer']
df = read_csv(model_path + '/' + 'df.csv')
scaler = GradScaler()
best_loss = 1e4
last_epoch_data_loss = 1.0
last_epoch_physical_loss = 1.0
for step in range(start_epoch, args.epoch):
print("epoch: ", step, flush=True)
reg_norm = regConstScheduler(step, args, last_epoch_data_loss,
last_epoch_physical_loss)
# training
for sample_batched in train_loader:
model.optimizer.zero_grad()
x_batch_train, y_batch_train = sample_batched['structure'].to(
device), sample_batched['field'].to(device)
with autocast():
logits = model(x_batch_train, bn_training=True)
#calculate the loss using the ground truth
loss = model.loss_fn(logits, y_batch_train)
logits = logits[:, :, 1:-1, :]
# print("loss: ", loss, flush=True)
# Calculate physical residue
pattern = (x_batch_train * (n_Si - n_air) + n_air)**2
# rescale the 0/1 pattern into dielectric constant
pattern = torch.cat(
(torch.ones([pattern.shape[0], 1, 1, 256],
dtype=torch.float32,
device=device) * n_sub**2, pattern),
dim=2)
#fields = logits; # predicted fields [Hy_R, Hy_I, Ex_R, Ex_I, Ez_R, Ez_I]
fields = torch.cat((y_batch_train[:, :, 0:1, :], logits,
y_batch_train[:, :, -1:, :]),
dim=2)
FD_Hy = H_to_H(-fields[:, 0] * means[0],
-fields[:, 1] * means[1], dL, omega, pattern)
#phys_regR = 10*model.loss_fn(FD_Hy[:, 0]/means[0], logits[:, 0])/reg_norm;
#phys_regI = 10*model.loss_fn(FD_Hy[:, 1]/means[1], logits[:, 1])/reg_norm;
phys_regR = model.loss_fn(FD_Hy[:, 0] / means[0],
logits[:, 0]) * reg_norm
phys_regI = model.loss_fn(FD_Hy[:, 1] / means[1],
logits[:, 1]) * reg_norm
loss += phys_regR + phys_regI
scaler.scale(loss).backward()
scaler.step(model.optimizer)
scaler.update()
#loss.backward()
#model.optimizer.step()
#Save the weights at the end of each epoch
checkpoint = {
'epoch': step,
'model': model,
'optimizer': model.optimizer,
'lr_scheduler': model.lr_scheduler
}
torch.save(checkpoint, model_path + "/last_model.pt")
# evaluation
train_loss = 0
train_phys_reg = 0
for sample_batched in train_loader:
x_batch_train, y_batch_train = sample_batched['structure'].to(
device), sample_batched['field'].to(device)
with torch.no_grad():
logits = model(x_batch_train, bn_training=False)
loss = model.loss_fn(logits, y_batch_train)
logits = logits[:, :, 1:-1, :]
# Calculate physical residue
pattern = (x_batch_train * (n_Si - n_air) + n_air)**2
# rescale the 0/1 pattern into dielectric constant
pattern = torch.cat(
(torch.ones([pattern.shape[0], 1, 1, 256],
dtype=torch.float32,
device=device) * n_sub**2, pattern),
dim=2)
#fields = logits; # predicted fields [Hy_R, Hy_I, Ex_R, Ex_I, Ez_R, Ez_I]
fields = torch.cat((y_batch_train[:, :, 0:1, :], logits,
y_batch_train[:, :, -1:, :]),
dim=2)
FD_Hy = H_to_H(-fields[:, 0] * means[0],
-fields[:, 1] * means[1], dL, omega, pattern)
#phys_regR = 10*model.loss_fn(FD_Hy[:, 0]/means[0], fields[:, 0, 1:-1, :])/reg_norm;
#phys_regI = 10*model.loss_fn(FD_Hy[:, 1]/means[1], fields[:, 1, 1:-1, :])/reg_norm;
phys_regR = model.loss_fn(FD_Hy[:, 0] / means[0],
fields[:, 0, 1:-1, :]) * reg_norm
phys_regI = model.loss_fn(FD_Hy[:, 1] / means[1],
fields[:, 1, 1:-1, :]) * reg_norm
#loss = loss + phys_reg1 + phys_reg2 + phys_reg3;
train_loss += loss
train_phys_reg += 0.5 * (phys_regR + phys_regI)
train_loss /= len(train_loader)
train_phys_reg /= len(train_loader)
test_loss = 0
test_phys_reg = 0
for sample_batched in test_loader:
x_batch_test, y_batch_test = sample_batched['structure'].to(
device), sample_batched['field'].to(device)
with torch.no_grad():
logits = model(x_batch_test, bn_training=False)
loss = model.loss_fn(logits, y_batch_test)
logits = logits[:, :, 1:-1, :]
# Calculate physical residue
pattern = (x_batch_test * (n_Si - n_air) + n_air)**2
# rescale the 0/1 pattern into dielectric constant
pattern = torch.cat(
(torch.ones([pattern.shape[0], 1, 1, 256],
dtype=torch.float32,
device=device) * n_sub**2, pattern),
dim=2)
#fields = logits; # predicted fields [Hy_R, Hy_I, Ex_R, Ex_I, Ez_R, Ez_I]
fields = torch.cat(
(y_batch_test[:, :, 0:1, :], logits, y_batch_test[:, :,
-1:, :]),
dim=2)
FD_Hy = H_to_H(-fields[:, 0] * means[0],
-fields[:, 1] * means[1], dL, omega, pattern)
#phys_regR = 10*model.loss_fn(FD_Hy[:, 0]/means[0], fields[:, 0, 1:-1, :])/reg_norm;
#phys_regI = 10*model.loss_fn(FD_Hy[:, 1]/means[1], fields[:, 1, 1:-1, :])/reg_norm;
phys_regR = model.loss_fn(FD_Hy[:, 0] / means[0],
fields[:, 0, 1:-1, :])
phys_regI = model.loss_fn(FD_Hy[:, 1] / means[1],
fields[:, 1, 1:-1, :])
test_loss += loss
test_phys_reg += 0.5 * (phys_regR + phys_regI)
test_loss /= len(test_loader)
test_phys_reg /= len(test_loader)
last_epoch_data_loss = test_loss
last_epoch_physical_loss = test_phys_reg.detach().clone()
test_phys_reg *= reg_norm
print(
'train loss: %.5f, train phys reg: %.5f, test loss: %.5f, test phys reg: %.5f, last_physical_loss: %.5f'
% (train_loss, train_phys_reg, test_loss, test_phys_reg,
last_epoch_physical_loss),
flush=True)
model.lr_scheduler.step()
df = df.append(
{
'epoch': step + 1,
'lr': str(model.lr_scheduler.get_last_lr()),
'train_loss': train_loss.item(),
'train_phys_reg': train_phys_reg.item(),
'test_loss': test_loss.item(),
'test_phys_reg': test_phys_reg.item(),
},
ignore_index=True)
df.to_csv(model_path + '/' + 'df.csv', index=False)
if (test_loss < best_loss):
best_loss = test_loss
checkpoint = {
'epoch': step,
'model': model,
'optimizer': model.optimizer,
'lr_scheduler': model.lr_scheduler
}
torch.save(checkpoint, model_path + "/best_model.pt")
class NetworkTrainer(object):
def __init__(self, deterministic=True, fp16=False):
"""
A generic class that can train almost any neural network (RNNs excluded). It provides basic functionality such
as the training loop, tracking of training and validation losses (and the target metric if you implement it)
Training can be terminated early if the validation loss (or the target metric if implemented) do not improve
anymore. This is based on a moving average (MA) of the loss/metric instead of the raw values to get more smooth
results.
What you need to override:
- __init__
- initialize
- run_online_evaluation (optional)
- finish_online_evaluation (optional)
- validate
- predict_test_case
"""
self.fp16 = fp16
self.amp_grad_scaler = None
if deterministic:
np.random.seed(12345)
torch.manual_seed(12345)
if torch.cuda.is_available():
torch.cuda.manual_seed_all(12345)
cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
else:
cudnn.deterministic = False
torch.backends.cudnn.benchmark = True
################# SET THESE IN self.initialize() ###################################
self.network: Tuple[SegmentationNetwork, nn.DataParallel] = None
self.optimizer = None
self.lr_scheduler = None
self.tr_gen = self.val_gen = None
self.was_initialized = False
################# SET THESE IN INIT ################################################
self.output_folder = None
self.fold = None
self.loss = None
self.dataset_directory = None
################# SET THESE IN LOAD_DATASET OR DO_SPLIT ############################
self.dataset = None # these can be None for inference mode
self.dataset_tr = self.dataset_val = None # do not need to be used, they just appear if you are using the suggested load_dataset_and_do_split
################# THESE DO NOT NECESSARILY NEED TO BE MODIFIED #####################
self.patience = 50
self.val_eval_criterion_alpha = 0.9 # alpha * old + (1-alpha) * new
# if this is too low then the moving average will be too noisy and the training may terminate early. If it is
# too high the training will take forever
self.train_loss_MA_alpha = 0.93 # alpha * old + (1-alpha) * new
self.train_loss_MA_eps = 5e-4 # new MA must be at least this much better (smaller)
self.max_num_epochs = 1000
self.num_batches_per_epoch = 250
self.num_val_batches_per_epoch = 50
self.also_val_in_tr_mode = False
self.lr_threshold = 1e-6 # the network will not terminate training if the lr is still above this threshold
################# LEAVE THESE ALONE ################################################
self.val_eval_criterion_MA = None
self.train_loss_MA = None
self.best_val_eval_criterion_MA = None
self.best_MA_tr_loss_for_patience = None
self.best_epoch_based_on_MA_tr_loss = None
self.all_tr_losses = []
self.all_val_losses = []
self.all_val_losses_tr_mode = []
self.all_val_eval_metrics = [] # does not have to be used
self.epoch = 0
self.log_file = None
self.deterministic = deterministic
self.use_progress_bar = False
if 'nnunet_use_progress_bar' in os.environ.keys():
self.use_progress_bar = bool(
int(os.environ['nnunet_use_progress_bar']))
################# Settings for saving checkpoints ##################################
self.save_every = 50
self.save_latest_only = True # if false it will not store/overwrite _latest but separate files each
# time an intermediate checkpoint is created
self.save_intermediate_checkpoints = True # whether or not to save checkpoint_latest
self.save_best_checkpoint = True # whether or not to save the best checkpoint according to self.best_val_eval_criterion_MA
self.save_final_checkpoint = True # whether or not to save the final checkpoint
@abstractmethod
def initialize(self, training=True):
"""
create self.output_folder
modify self.output_folder if you are doing cross-validation (one folder per fold)
set self.tr_gen and self.val_gen
call self.initialize_network and self.initialize_optimizer_and_scheduler (important!)
finally set self.was_initialized to True
:param training:
:return:
"""
@abstractmethod
def load_dataset(self):
pass
def do_split(self):
"""
This is a suggestion for if your dataset is a dictionary (my personal standard)
:return:
"""
splits_file = join(self.dataset_directory, "splits_final.pkl")
if not isfile(splits_file):
self.print_to_log_file("Creating new split...")
splits = []
all_keys_sorted = np.sort(list(self.dataset.keys()))
kfold = KFold(n_splits=5, shuffle=True, random_state=12345)
for i, (train_idx,
test_idx) in enumerate(kfold.split(all_keys_sorted)):
train_keys = np.array(all_keys_sorted)[train_idx]
test_keys = np.array(all_keys_sorted)[test_idx]
splits.append(OrderedDict())
splits[-1]['train'] = train_keys
splits[-1]['val'] = test_keys
save_pickle(splits, splits_file)
splits = load_pickle(splits_file)
if self.fold == "all":
tr_keys = val_keys = list(self.dataset.keys())
else:
tr_keys = splits[self.fold]['train']
val_keys = splits[self.fold]['val']
tr_keys.sort()
val_keys.sort()
self.dataset_tr = OrderedDict()
for i in tr_keys:
self.dataset_tr[i] = self.dataset[i]
self.dataset_val = OrderedDict()
for i in val_keys:
self.dataset_val[i] = self.dataset[i]
def plot_progress(self):
"""
Should probably by improved
:return:
"""
try:
font = {'weight': 'normal', 'size': 18}
matplotlib.rc('font', **font)
fig = plt.figure(figsize=(30, 24))
ax = fig.add_subplot(111)
ax2 = ax.twinx()
x_values = list(range(self.epoch + 1))
ax.plot(x_values,
self.all_tr_losses,
color='b',
ls='-',
label="loss_tr")
ax.plot(x_values,
self.all_val_losses,
color='r',
ls='-',
label="loss_val, train=False")
if len(self.all_val_losses_tr_mode) > 0:
ax.plot(x_values,
self.all_val_losses_tr_mode,
color='g',
ls='-',
label="loss_val, train=True")
if len(self.all_val_eval_metrics) == len(x_values):
ax2.plot(x_values,
self.all_val_eval_metrics,
color='g',
ls='--',
label="evaluation metric")
ax.set_xlabel("epoch")
ax.set_ylabel("loss")
ax2.set_ylabel("evaluation metric")
ax.legend()
ax2.legend(loc=9)
fig.savefig(join(self.output_folder, "progress.png"))
plt.close()
except IOError:
self.print_to_log_file("failed to plot: ", sys.exc_info())
def print_to_log_file(self,
*args,
also_print_to_console=True,
add_timestamp=True):
timestamp = time()
dt_object = datetime.fromtimestamp(timestamp)
if add_timestamp:
args = ("%s:" % dt_object, *args)
if self.log_file is None:
maybe_mkdir_p(self.output_folder)
timestamp = datetime.now()
self.log_file = join(
self.output_folder,
"training_log_%d_%d_%d_%02.0d_%02.0d_%02.0d.txt" %
(timestamp.year, timestamp.month, timestamp.day,
timestamp.hour, timestamp.minute, timestamp.second))
with open(self.log_file, 'w') as f:
f.write("Starting... \n")
successful = False
max_attempts = 5
ctr = 0
while not successful and ctr < max_attempts:
try:
with open(self.log_file, 'a+') as f:
for a in args:
f.write(str(a))
f.write(" ")
f.write("\n")
successful = True
except IOError:
print(
"%s: failed to log: " % datetime.fromtimestamp(timestamp),
sys.exc_info())
sleep(0.5)
ctr += 1
if also_print_to_console:
print(*args)
def save_checkpoint(self, fname, save_optimizer=True):
start_time = time()
state_dict = self.network.state_dict()
for key in state_dict.keys():
state_dict[key] = state_dict[key].cpu()
lr_sched_state_dct = None
if self.lr_scheduler is not None and hasattr(
self.lr_scheduler, 'state_dict'
): # not isinstance(self.lr_scheduler, lr_scheduler.ReduceLROnPlateau):
lr_sched_state_dct = self.lr_scheduler.state_dict()
# WTF is this!?
# for key in lr_sched_state_dct.keys():
# lr_sched_state_dct[key] = lr_sched_state_dct[key]
if save_optimizer:
optimizer_state_dict = self.optimizer.state_dict()
else:
optimizer_state_dict = None
self.print_to_log_file("saving checkpoint...")
save_this = {
'epoch':
self.epoch + 1,
'state_dict':
state_dict,
'optimizer_state_dict':
optimizer_state_dict,
'lr_scheduler_state_dict':
lr_sched_state_dct,
'plot_stuff':
(self.all_tr_losses, self.all_val_losses,
self.all_val_losses_tr_mode, self.all_val_eval_metrics),
'best_stuff': (self.best_epoch_based_on_MA_tr_loss,
self.best_MA_tr_loss_for_patience,
self.best_val_eval_criterion_MA)
}
if self.amp_grad_scaler is not None:
save_this['amp_grad_scaler'] = self.amp_grad_scaler.state_dict()
torch.save(save_this, fname)
self.print_to_log_file("done, saving took %.2f seconds" %
(time() - start_time))
def load_best_checkpoint(self, train=True):
if self.fold is None:
raise RuntimeError(
"Cannot load best checkpoint if self.fold is None")
if isfile(join(self.output_folder, "model_best.model")):
self.load_checkpoint(join(self.output_folder, "model_best.model"),
train=train)
else:
self.print_to_log_file(
"WARNING! model_best.model does not exist! Cannot load best checkpoint. Falling "
"back to load_latest_checkpoint")
self.load_latest_checkpoint(train)
def load_latest_checkpoint(self, train=True):
if isfile(join(self.output_folder, "model_final_checkpoint.model")):
return self.load_checkpoint(join(self.output_folder,
"model_final_checkpoint.model"),
train=train)
if isfile(join(self.output_folder, "model_latest.model")):
return self.load_checkpoint(join(self.output_folder,
"model_latest.model"),
train=train)
if isfile(join(self.output_folder, "model_best.model")):
return self.load_best_checkpoint(train)
raise RuntimeError("No checkpoint found")
def load_checkpoint(self, fname, train=True):
self.print_to_log_file("loading checkpoint", fname, "train=", train)
if not self.was_initialized:
self.initialize(train)
# saved_model = torch.load(fname, map_location=torch.device('cuda', torch.cuda.current_device()))
saved_model = torch.load(fname, map_location=torch.device('cpu'))
self.load_checkpoint_ram(saved_model, train)
@abstractmethod
def initialize_network(self):
"""
initialize self.network here
:return:
"""
pass
@abstractmethod
def initialize_optimizer_and_scheduler(self):
"""
initialize self.optimizer and self.lr_scheduler (if applicable) here
:return:
"""
pass
def load_checkpoint_ram(self, checkpoint, train=True):
"""
used for if the checkpoint is already in ram
:param checkpoint:
:param train:
:return:
"""
if not self.was_initialized:
self.initialize(train)
new_state_dict = OrderedDict()
curr_state_dict_keys = list(self.network.state_dict().keys())
# if state dict comes form nn.DataParallel but we use non-parallel model here then the state dict keys do not
# match. Use heuristic to make it match
for k, value in checkpoint['state_dict'].items():
key = k
if key not in curr_state_dict_keys and key.startswith('module.'):
key = key[7:]
new_state_dict[key] = value
if self.fp16:
self._maybe_init_amp()
if 'amp_grad_scaler' in checkpoint.keys():
self.amp_grad_scaler.load_state_dict(
checkpoint['amp_grad_scaler'])
self.network.load_state_dict(new_state_dict)
self.epoch = checkpoint['epoch']
if train:
optimizer_state_dict = checkpoint['optimizer_state_dict']
if optimizer_state_dict is not None:
self.optimizer.load_state_dict(optimizer_state_dict)
if self.lr_scheduler is not None and hasattr(
self.lr_scheduler, 'load_state_dict'
) and checkpoint['lr_scheduler_state_dict'] is not None:
self.lr_scheduler.load_state_dict(
checkpoint['lr_scheduler_state_dict'])
if issubclass(self.lr_scheduler.__class__, _LRScheduler):
self.lr_scheduler.step(self.epoch)
self.all_tr_losses, self.all_val_losses, self.all_val_losses_tr_mode, self.all_val_eval_metrics = checkpoint[
'plot_stuff']
# load best loss (if present)
if 'best_stuff' in checkpoint.keys():
self.best_epoch_based_on_MA_tr_loss, self.best_MA_tr_loss_for_patience, self.best_val_eval_criterion_MA = checkpoint[
'best_stuff']
# after the training is done, the epoch is incremented one more time in my old code. This results in
# self.epoch = 1001 for old trained models when the epoch is actually 1000. This causes issues because
# len(self.all_tr_losses) = 1000 and the plot function will fail. We can easily detect and correct that here
if self.epoch != len(self.all_tr_losses):
self.print_to_log_file(
"WARNING in loading checkpoint: self.epoch != len(self.all_tr_losses). This is "
"due to an old bug and should only appear when you are loading old models. New "
"models should have this fixed! self.epoch is now set to len(self.all_tr_losses)"
)
self.epoch = len(self.all_tr_losses)
self.all_tr_losses = self.all_tr_losses[:self.epoch]
self.all_val_losses = self.all_val_losses[:self.epoch]
self.all_val_losses_tr_mode = self.all_val_losses_tr_mode[:self.
epoch]
self.all_val_eval_metrics = self.all_val_eval_metrics[:self.epoch]
self._maybe_init_amp()
def _maybe_init_amp(self):
if self.fp16 and self.amp_grad_scaler is None and torch.cuda.is_available(
):
self.amp_grad_scaler = GradScaler()
def plot_network_architecture(self):
"""
can be implemented (see nnUNetTrainer) but does not have to. Not implemented here because it imposes stronger
assumptions on the presence of class variables
:return:
"""
pass
def run_training(self):
_ = self.tr_gen.next()
_ = self.val_gen.next()
if torch.cuda.is_available():
torch.cuda.empty_cache()
self._maybe_init_amp()
maybe_mkdir_p(self.output_folder)
self.plot_network_architecture()
if cudnn.benchmark and cudnn.deterministic:
warn(
"torch.backends.cudnn.deterministic is True indicating a deterministic training is desired. "
"But torch.backends.cudnn.benchmark is True as well and this will prevent deterministic training! "
"If you want deterministic then set benchmark=False")
if not self.was_initialized:
self.initialize(True)
while self.epoch < self.max_num_epochs:
self.print_to_log_file("\nepoch: ", self.epoch)
epoch_start_time = time()
train_losses_epoch = []
# train one epoch
self.network.train()
if self.use_progress_bar:
with trange(self.num_batches_per_epoch) as tbar:
for b in tbar:
tbar.set_description("Epoch {}/{}".format(
self.epoch + 1, self.max_num_epochs))
l = self.run_iteration(self.tr_gen, True)
tbar.set_postfix(loss=l)
train_losses_epoch.append(l)
else:
for _ in range(self.num_batches_per_epoch):
l = self.run_iteration(self.tr_gen, True)
train_losses_epoch.append(l)
self.all_tr_losses.append(np.mean(train_losses_epoch))
self.print_to_log_file("train loss : %.4f" %
self.all_tr_losses[-1])
with torch.no_grad():
# validation with train=False
self.network.eval()
val_losses = []
for b in range(self.num_val_batches_per_epoch):
l = self.run_iteration(self.val_gen, False, True)
val_losses.append(l)
self.all_val_losses.append(np.mean(val_losses))
self.print_to_log_file("validation loss: %.4f" %
self.all_val_losses[-1])
if self.also_val_in_tr_mode:
self.network.train()
# validation with train=True
val_losses = []
for b in range(self.num_val_batches_per_epoch):
l = self.run_iteration(self.val_gen, False)
val_losses.append(l)
self.all_val_losses_tr_mode.append(np.mean(val_losses))
self.print_to_log_file(
"validation loss (train=True): %.4f" %
self.all_val_losses_tr_mode[-1])
self.update_train_loss_MA(
) # needed for lr scheduler and stopping of training
continue_training = self.on_epoch_end()
epoch_end_time = time()
if not continue_training:
# allows for early stopping
break
self.epoch += 1
self.print_to_log_file("This epoch took %f s\n" %
(epoch_end_time - epoch_start_time))
self.epoch -= 1 # if we don't do this we can get a problem with loading model_final_checkpoint.
if self.save_final_checkpoint:
self.save_checkpoint(
join(self.output_folder, "model_final_checkpoint.model"))
# now we can delete latest as it will be identical with final
if isfile(join(self.output_folder, "model_latest.model")):
os.remove(join(self.output_folder, "model_latest.model"))
if isfile(join(self.output_folder, "model_latest.model.pkl")):
os.remove(join(self.output_folder, "model_latest.model.pkl"))
def maybe_update_lr(self):
# maybe update learning rate
if self.lr_scheduler is not None:
assert isinstance(
self.lr_scheduler,
(lr_scheduler.ReduceLROnPlateau, lr_scheduler._LRScheduler))
if isinstance(self.lr_scheduler, lr_scheduler.ReduceLROnPlateau):
# lr scheduler is updated with moving average val loss. should be more robust
self.lr_scheduler.step(self.train_loss_MA)
else:
self.lr_scheduler.step(self.epoch + 1)
self.print_to_log_file("lr is now (scheduler) %s" %
str(self.optimizer.param_groups[0]['lr']))
def maybe_save_checkpoint(self):
"""
Saves a checkpoint every save_ever epochs.
:return:
"""
if self.save_intermediate_checkpoints and (self.epoch % self.save_every
== (self.save_every - 1)):
self.print_to_log_file("saving scheduled checkpoint file...")
if not self.save_latest_only:
self.save_checkpoint(
join(self.output_folder,
"model_ep_%03.0d.model" % (self.epoch + 1)))
self.save_checkpoint(join(self.output_folder,
"model_latest.model"))
self.print_to_log_file("done")
def update_eval_criterion_MA(self):
"""
If self.all_val_eval_metrics is unused (len=0) then we fall back to using -self.all_val_losses for the MA to determine early stopping
(not a minimization, but a maximization of a metric and therefore the - in the latter case)
:return:
"""
if self.val_eval_criterion_MA is None:
if len(self.all_val_eval_metrics) == 0:
self.val_eval_criterion_MA = -self.all_val_losses[-1]
else:
self.val_eval_criterion_MA = self.all_val_eval_metrics[-1]
else:
if len(self.all_val_eval_metrics) == 0:
"""
We here use alpha * old - (1 - alpha) * new because new in this case is the vlaidation loss and lower
is better, so we need to negate it.
"""
self.val_eval_criterion_MA = self.val_eval_criterion_alpha * self.val_eval_criterion_MA - (
1 - self.val_eval_criterion_alpha) * \
self.all_val_losses[-1]
else:
self.val_eval_criterion_MA = self.val_eval_criterion_alpha * self.val_eval_criterion_MA + (
1 - self.val_eval_criterion_alpha) * \
self.all_val_eval_metrics[-1]
def manage_patience(self):
# update patience
continue_training = True
if self.patience is not None:
# if best_MA_tr_loss_for_patience and best_epoch_based_on_MA_tr_loss were not yet initialized,
# initialize them
if self.best_MA_tr_loss_for_patience is None:
self.best_MA_tr_loss_for_patience = self.train_loss_MA
if self.best_epoch_based_on_MA_tr_loss is None:
self.best_epoch_based_on_MA_tr_loss = self.epoch
if self.best_val_eval_criterion_MA is None:
self.best_val_eval_criterion_MA = self.val_eval_criterion_MA
# check if the current epoch is the best one according to moving average of validation criterion. If so
# then save 'best' model
# Do not use this for validation. This is intended for test set prediction only.
#self.print_to_log_file("current best_val_eval_criterion_MA is %.4f0" % self.best_val_eval_criterion_MA)
#self.print_to_log_file("current val_eval_criterion_MA is %.4f" % self.val_eval_criterion_MA)
if self.val_eval_criterion_MA > self.best_val_eval_criterion_MA:
self.best_val_eval_criterion_MA = self.val_eval_criterion_MA
#self.print_to_log_file("saving best epoch checkpoint...")
if self.save_best_checkpoint:
self.save_checkpoint(
join(self.output_folder, "model_best.model"))
# Now see if the moving average of the train loss has improved. If yes then reset patience, else
# increase patience
if self.train_loss_MA + self.train_loss_MA_eps < self.best_MA_tr_loss_for_patience:
self.best_MA_tr_loss_for_patience = self.train_loss_MA
self.best_epoch_based_on_MA_tr_loss = self.epoch
#self.print_to_log_file("New best epoch (train loss MA): %03.4f" % self.best_MA_tr_loss_for_patience)
else:
pass
#self.print_to_log_file("No improvement: current train MA %03.4f, best: %03.4f, eps is %03.4f" %
# (self.train_loss_MA, self.best_MA_tr_loss_for_patience, self.train_loss_MA_eps))
# if patience has reached its maximum then finish training (provided lr is low enough)
if self.epoch - self.best_epoch_based_on_MA_tr_loss > self.patience:
if self.optimizer.param_groups[0]['lr'] > self.lr_threshold:
#self.print_to_log_file("My patience ended, but I believe I need more time (lr > 1e-6)")
self.best_epoch_based_on_MA_tr_loss = self.epoch - self.patience // 2
else:
#self.print_to_log_file("My patience ended")
continue_training = False
else:
pass
#self.print_to_log_file(
# "Patience: %d/%d" % (self.epoch - self.best_epoch_based_on_MA_tr_loss, self.patience))
return continue_training
def on_epoch_end(self):
self.finish_online_evaluation(
) # does not have to do anything, but can be used to update self.all_val_eval_
# metrics
self.plot_progress()
self.maybe_update_lr()
self.maybe_save_checkpoint()
self.update_eval_criterion_MA()
continue_training = self.manage_patience()
return continue_training
def update_train_loss_MA(self):
if self.train_loss_MA is None:
self.train_loss_MA = self.all_tr_losses[-1]
else:
self.train_loss_MA = self.train_loss_MA_alpha * self.train_loss_MA + (1 - self.train_loss_MA_alpha) * \
self.all_tr_losses[-1]
def run_iteration(self,
data_generator,
do_backprop=True,
run_online_evaluation=False):
data_dict = next(data_generator)
data = data_dict['data']
target = data_dict['target']
data = maybe_to_torch(data)
target = maybe_to_torch(target)
if torch.cuda.is_available():
data = to_cuda(data)
target = to_cuda(target)
self.optimizer.zero_grad()
if self.fp16:
with autocast():
output = self.network(data)
del data
l = self.loss(output, target)
if do_backprop:
self.amp_grad_scaler.scale(l).backward()
self.amp_grad_scaler.step(self.optimizer)
self.amp_grad_scaler.update()
else:
output = self.network(data)
del data
l = self.loss(output, target)
if do_backprop:
l.backward()
self.optimizer.step()
if run_online_evaluation:
self.run_online_evaluation(output, target)
del target
return l.detach().cpu().numpy()
def run_online_evaluation(self, *args, **kwargs):
"""
Can be implemented, does not have to
:param output_torch:
:param target_npy:
:return:
"""
pass
def finish_online_evaluation(self):
"""
Can be implemented, does not have to
:return:
"""
pass
@abstractmethod
def validate(self, *args, **kwargs):
pass
def find_lr(self,
num_iters=1000,
init_value=1e-6,
final_value=10.,
beta=0.98):
"""
stolen and adapted from here: https://sgugger.github.io/how-do-you-find-a-good-learning-rate.html
:param num_iters:
:param init_value:
:param final_value:
:param beta:
:return:
"""
import math
self._maybe_init_amp()
mult = (final_value / init_value)**(1 / num_iters)
lr = init_value
self.optimizer.param_groups[0]['lr'] = lr
avg_loss = 0.
best_loss = 0.
losses = []
log_lrs = []
for batch_num in range(1, num_iters + 1):
# +1 because this one here is not designed to have negative loss...
loss = self.run_iteration(
self.tr_gen, do_backprop=True,
run_online_evaluation=False).data.item() + 1
# Compute the smoothed loss
avg_loss = beta * avg_loss + (1 - beta) * loss
smoothed_loss = avg_loss / (1 - beta**batch_num)
# Stop if the loss is exploding
if batch_num > 1 and smoothed_loss > 4 * best_loss:
break
# Record the best loss
if smoothed_loss < best_loss or batch_num == 1:
best_loss = smoothed_loss
# Store the values
losses.append(smoothed_loss)
log_lrs.append(math.log10(lr))
# Update the lr for the next step
lr *= mult
self.optimizer.param_groups[0]['lr'] = lr
import matplotlib.pyplot as plt
lrs = [10**i for i in log_lrs]
fig = plt.figure()
plt.xscale('log')
plt.plot(lrs[10:-5], losses[10:-5])
plt.savefig(join(self.output_folder, "lr_finder.png"))
plt.close()
return log_lrs, losses
class ClassificationTask(ClassyTask):
"""Basic classification training task.
This task encapsultates all of the components and steps needed to
train a classifier using a :class:`classy_vision.trainer.ClassyTrainer`.
Assumes a train / test phase per each epoch and that the datasets
have the same API as the map-style Dataset class in
`torch.utils.data.dataset <https://pytorch.org/docs/stable/data.html
#torch.utils.data.Dataset>`_ (in particular, this task makes use of
the len). If you are using an `IterableDataset <https://pytorch.org/docs/
stable/data.html#torch.utils.data.IterableDataset>`_ then a custom task
may be appropriate.
:var loss: Loss (see :class:`classy_vision.losses.ClassyLoss`) function used
for computing the loss in each forward pass
:var datasets: Mapping from a ``phase_type`` in ["train", "test']
to dataset used for training (or testing)
:var meters: List of meters (see :class:`classy_vision.meters.ClassyMeter`)
to calculate during training
:var num_epochs: Number of epochs (passes over dataset) to train
:var test_only: Used to only run the test phase
:var base_model: Model to be trained, unwrapped in DDP or DP wrappers
:var optimizer: Optimizer used in train step
:var optimizer_schedulers: Dictionary. Key is the name of the optimizer
option (e.g. lr), value is a ClassyParamScheduler
:var checkpoint: Serializable dict which represents state in training
:var phases: List of phase specific information, e.g. if phase is
train / test.
:var hooks: List of hooks to apply during training
:var train: Phase type, if true it means we are training,
false means testing
:var distributed_model: Base model, but wrapped in DDP (DistributedDataParallel)
:var phase_idx: Current phase id, first phase is 0, if task has not started
training then returns -1
:var train_phase_idx: Only counts train phases
:var num_updates: Number of total parameter updates applied to model
by the optimizer
:var data_iterator: Iterator which can be used to obtain batches
:var losses: Loss curve
:var perf_log: list of training speed measurements, to be logged
:var clip_grad_norm: maximum gradient norm (default None)
:var simulated_global_batchsize: batch size simulated via gradient accumulation
:var optimizer_period: apply optimizer after this many steps; derived from
simulated_global_batchsize, default 1.
"""
def __init__(self):
"""Constructs a ClassificationTask"""
super().__init__()
self.base_loss = None
self.datasets = {}
self.meters = []
self.num_epochs = 1
self.test_phase_period = 1
self.train_phases_per_epoch = 0
self.test_only = False
self.base_model = None
self.optimizer = None
self.optimizer_schedulers = {}
self.checkpoint_dict = None
self.checkpoint_path = None
self.phases = []
self.hooks = []
self.train = True
self.distributed_model = None
self.distributed_loss = None
self.phase_idx = -1
self.train_phase_idx = -1
self.num_updates = 0
self.dataloader = None
self.data_iterator = None
self.losses = []
self.broadcast_buffers_mode: BroadcastBuffersMode = (
BroadcastBuffersMode.BEFORE_EVAL
)
self.amp_args = None
self.amp_type = None
self.amp_grad_scaler = None
self.mixup_transform = None
self.perf_log = []
self.last_batch = None
self.batch_norm_sync_mode = BatchNormSyncMode.DISABLED
self.find_unused_parameters = False
self.use_gpu = torch.cuda.is_available()
self.dataloader_mp_context = "spawn"
self.bn_weight_decay = False
self._train_only = True
self.clip_grad_norm = None
self.simulated_global_batchsize = None
self.optimizer_period = 1
self.ddp_bucket_cap_mb = 25
self.use_sharded_ddp = False
self.fp16_grad_compress = False
def set_use_sharded_ddp(self, use_sharded_ddp: bool):
self.use_sharded_ddp = use_sharded_ddp
if self.use_sharded_ddp:
logging.info("Using Sharded DDP")
return self
def set_use_gpu(self, use_gpu: bool):
self.use_gpu = use_gpu
assert (
not self.use_gpu or torch.cuda.is_available()
), "CUDA required to train on GPUs"
return self
def set_clip_grad_norm(self, clip_grad_norm: Optional[float]):
"""Sets maximum gradient norm.
None means gradient clipping is disabled. Defaults to None."""
self.clip_grad_norm = clip_grad_norm
if clip_grad_norm is None:
logging.info("Disabled gradient norm clipping.")
else:
logging.info(
f"Enabled gradient norm clipping with threshold: {clip_grad_norm}"
)
return self
def set_simulated_global_batchsize(self, simulated_global_batchsize: Optional[int]):
"""Sets a simulated batch size by gradient accumulation.
Gradient accumulation adds up gradients from multiple minibatches and
steps the optimizer every N train_steps, where N is optimizer_period.
When enabled, the very last train_steps might end up not updating the
model, depending on the number of total steps. None means gradient
accumulation is disabled. Defaults to None."""
self.simulated_global_batchsize = simulated_global_batchsize
return self
def set_checkpoint(self, checkpoint_path: str):
"""Sets checkpoint on task.
Args:
checkpoint_path: The path to load the checkpoint from. Can be a file or a
directory. See :func:`load_checkpoint` for more information.
"""
self.checkpoint_path = checkpoint_path
return self
def _set_checkpoint_dict(self, checkpoint_dict: Dict[str, Any]):
"""Sets the checkpoint dict in the task. Only used for testing.
Args:
checkpoint_dict: A serializable dict representing current task state
"""
self.checkpoint_dict = checkpoint_dict
return self
def set_num_epochs(self, num_epochs: Union[int, float]):
"""Set number of epochs to be run.
Args:
num_epochs: Number of epochs to run task
"""
self.num_epochs = num_epochs
return self
def set_test_phase_period(self, test_phase_period: int):
"""Set the period of test phase.
Args:
test_phase_period: The period of test phase
"""
self.test_phase_period = test_phase_period
return self
def set_dataset(self, dataset: ClassyDataset, phase_type: str):
"""Set dataset for phase type on task
Args:
dataset: ClassyDataset for returning samples.
phase_type: str must be one of "train" or "test"
"""
assert phase_type in [
"train",
"test",
], "phase_type must be in ['train', 'test']"
self.datasets[phase_type] = dataset
if phase_type == "train":
self.train_phases_per_epoch = getattr(dataset, "phases_per_epoch", 1)
else:
self._train_only = False
return self
def set_dataloader_mp_context(self, dataloader_mp_context: Optional[str]):
"""Set the multiprocessing context used by the dataloader.
The context can be either 'spawn', 'fork', 'forkserver' or None (uses the
default context). See
https://docs.python.org/3/library/multiprocessing.html#multiprocessing.get_context
for more details."""
self.dataloader_mp_context = dataloader_mp_context
return self
def set_optimizer(self, optimizer: ClassyOptimizer):
"""Set optimizer for task
Args:
optimizer: optimizer for task
"""
self.optimizer = optimizer
return self
def set_loss(self, loss: ClassyLoss):
"""Set loss function for task
Args:
loss: loss for task
"""
self.base_loss = loss
return self
def set_meters(self, meters: List["ClassyMeter"]):
"""Set meters for task
Args:
meters: list of meters to compute during training
"""
self.meters = meters
return self
def set_distributed_options(
self,
broadcast_buffers_mode: BroadcastBuffersMode = BroadcastBuffersMode.BEFORE_EVAL,
batch_norm_sync_mode: BatchNormSyncMode = BatchNormSyncMode.DISABLED,
batch_norm_sync_group_size: int = 0,
find_unused_parameters: bool = False,
bucket_cap_mb: int = 25,
fp16_grad_compress: bool = False,
):
"""Set distributed options.
Args:
broadcast_buffers_mode: Broadcast buffers mode. See
:class:`BroadcastBuffersMode` for options.
batch_norm_sync_mode: Batch normalization synchronization mode. See
:class:`BatchNormSyncMode` for options.
batch_norm_sync_group_size: Group size to use for synchronized batch norm.
0 means that the stats are synchronized across all replicas. For
efficient synchronization, set it to the number of GPUs in a node (
usually 8).
find_unused_parameters: See
:class:`torch.nn.parallel.DistributedDataParallel` for information.
bucket_cap_mb: See
:class:`torch.nn.parallel.DistributedDataParallel` for information.
Raises:
RuntimeError: If batch_norm_sync_mode is `BatchNormSyncMode.APEX` and apex
is not installed.
"""
self.broadcast_buffers_mode = broadcast_buffers_mode
if batch_norm_sync_group_size > 0:
if not batch_norm_sync_mode == BatchNormSyncMode.APEX:
# this should ideally work with PyTorch Sync BN as well, but it
# fails while initializing DDP for some reason.
raise ValueError(
"batch_norm_sync_group_size can be > 0 only when "
"Apex Synchronized Batch Normalization is being used."
)
self.batch_norm_sync_group_size = batch_norm_sync_group_size
if batch_norm_sync_mode == BatchNormSyncMode.DISABLED:
logging.info("Synchronized Batch Normalization is disabled")
else:
if batch_norm_sync_mode == BatchNormSyncMode.APEX and not apex_available:
raise RuntimeError("apex is not installed")
msg = f"Using Synchronized Batch Normalization using {batch_norm_sync_mode}"
if self.batch_norm_sync_group_size > 0:
msg += f" and group size {batch_norm_sync_group_size}"
logging.info(msg)
self.batch_norm_sync_mode = batch_norm_sync_mode
if find_unused_parameters:
logging.info("Enabling find_unused_parameters in DDP")
self.find_unused_parameters = find_unused_parameters
self.ddp_bucket_cap_mb = bucket_cap_mb
if fp16_grad_compress:
if get_torch_version() < [1, 8]:
raise RuntimeError(
"FP16 grad compression is only supported since PyTorch 1.8"
)
logging.info("Enabling FP16 grad compression")
self.fp16_grad_compress = fp16_grad_compress
return self
def set_hooks(self, hooks: List["ClassyHook"]):
"""Set hooks for task
Args:
hooks: List of hooks to apply during training
"""
from classy_vision.hooks import ClassyHook
assert isinstance(hooks, list)
assert all(isinstance(hook, ClassyHook) for hook in hooks)
assert len({hook.name() for hook in hooks}) == len(
hooks
), "Cannot have repeated hooks of the same class"
# TODO (zyan3): we move checkpoint hook to the end of the list because some hooks
# may change the state of the model, and we want to save changed state in the checkpoint.
# This is temporary fix.
non_checkpoint_hooks = [
hook for hook in hooks if not isinstance(hook, CheckpointHook)
]
checkpoint_hooks = [hook for hook in hooks if isinstance(hook, CheckpointHook)]
hooks = non_checkpoint_hooks + checkpoint_hooks
self.hooks = hooks
return self
def set_model(self, model: ClassyModel):
"""Set model for task
Args:
model: Model to be trained
"""
self.base_model = model
return self
def set_test_only(self, test_only: bool):
"""Set test only flag
Args:
test_only: If true, only test phases will be run
"""
self.test_only = test_only
return self
def set_bn_weight_decay(self, bn_weight_decay: bool):
assert type(bn_weight_decay) == bool
self.bn_weight_decay = bn_weight_decay
return self
def set_amp_args(self, amp_args: Optional[Dict[str, Any]]):
"""Disable / enable apex.amp and set the automatic mixed precision parameters.
apex.amp can be utilized for mixed / half precision training.
Args:
amp_args: Dictionary containing arguments to be passed to
amp.initialize. Set to None to disable amp. To enable mixed
precision training, pass amp_args={"opt_level": "O1"} here.
See https://nvidia.github.io/apex/amp.html for more info.
Raises:
RuntimeError: If opt_level is not None and apex is not installed.
Warning: apex needs to be installed to utilize this feature.
"""
self.amp_args = amp_args
if amp_args is None:
logging.info("AMP disabled")
else:
# Check that the requested AMP type is known
try:
self.amp_type = AmpType[self.amp_args["amp_type"].upper()]
except KeyError:
logging.info("AMP type not specified, defaulting to Apex")
self.amp_type = AmpType.APEX
# Check for CUDA availability, required for both Apex and Pytorch AMP
if not torch.cuda.is_available():
raise RuntimeError(
"AMP is required but CUDA is not supported, cannot enable AMP"
)
# Check for Apex availability
if self.amp_type == AmpType.APEX and not apex_available:
raise RuntimeError(
"Apex AMP is required but Apex is not installed, cannot enable AMP"
)
if self.use_sharded_ddp:
if self.amp_type == AmpType.APEX:
raise RuntimeError(
"ShardedDDP has been requested, which is incompatible with Apex AMP"
)
if not fairscale_available:
raise RuntimeError(
"ShardedDDP has been requested, but fairscale is not installed in the current environment"
)
# Set Torch AMP grad scaler, used to prevent gradient underflow
elif self.amp_type == AmpType.PYTORCH:
if self.use_sharded_ddp:
logging.info("Using ShardedGradScaler to manage Pytorch AMP")
self.amp_grad_scaler = ShardedGradScaler()
else:
self.amp_grad_scaler = TorchGradScaler()
logging.info(f"AMP enabled with args {amp_args}")
return self
def set_mixup_transform(self, mixup_transform: Optional["MixupTransform"]):
"""Disable / enable mixup transform for data augmentation
Args::
mixup_transform: a callable object which performs mixup data augmentation
"""
self.mixup_transform = mixup_transform
if mixup_transform is None:
logging.info("mixup disabled")
else:
logging.info("mixup enabled")
return self
def set_optimizer_schedulers(self, schedulers):
self.optimizer_schedulers = schedulers
return self
@classmethod
def from_config(cls, config: Dict[str, Any]) -> "ClassificationTask":
"""Instantiates a ClassificationTask from a configuration.
Args:
config: A configuration for a ClassificationTask.
See :func:`__init__` for parameters expected in the config.
Returns:
A ClassificationTask instance.
"""
test_only = config.get("test_only", False)
if not test_only:
# TODO Make distinction between epochs and phases in optimizer clear
train_phases_per_epoch = config["dataset"]["train"].get(
"phases_per_epoch", 1
)
optimizer_config = config["optimizer"]
optimizer_config["num_epochs"] = (
config["num_epochs"] * train_phases_per_epoch
)
optimizer = build_optimizer(optimizer_config)
param_schedulers = build_optimizer_schedulers(optimizer_config)
datasets = {}
phase_types = ["train", "test"]
for phase_type in phase_types:
if phase_type in config["dataset"]:
datasets[phase_type] = build_dataset(config["dataset"][phase_type])
loss = build_loss(config["loss"])
amp_args = config.get("amp_args")
meters = build_meters(config.get("meters", {}))
model = build_model(config["model"])
mixup_transform = None
if config.get("mixup") is not None:
assert "alpha" in config["mixup"], "key alpha is missing in mixup dict"
mixup_transform = MixupTransform(
config["mixup"]["alpha"], config["mixup"].get("num_classes")
)
# hooks config is optional
hooks_config = config.get("hooks")
hooks = []
if hooks_config is not None:
hooks = build_hooks(hooks_config)
distributed_config = config.get("distributed", {})
distributed_options = {
"broadcast_buffers_mode": BroadcastBuffersMode[
distributed_config.get("broadcast_buffers", "before_eval").upper()
],
"batch_norm_sync_mode": BatchNormSyncMode[
distributed_config.get("batch_norm_sync_mode", "disabled").upper()
],
"batch_norm_sync_group_size": distributed_config.get(
"batch_norm_sync_group_size", 0
),
"find_unused_parameters": distributed_config.get(
"find_unused_parameters", False
),
"bucket_cap_mb": distributed_config.get("bucket_cap_mb", 25),
"fp16_grad_compress": distributed_config.get("fp16_grad_compress", False),
}
task = (
cls()
.set_num_epochs(config["num_epochs"])
.set_test_phase_period(config.get("test_phase_period", 1))
.set_loss(loss)
.set_test_only(test_only)
.set_model(model)
.set_meters(meters)
.set_amp_args(amp_args)
.set_mixup_transform(mixup_transform)
.set_distributed_options(**distributed_options)
.set_hooks(hooks)
.set_bn_weight_decay(config.get("bn_weight_decay", False))
.set_clip_grad_norm(config.get("clip_grad_norm"))
.set_simulated_global_batchsize(config.get("simulated_global_batchsize"))
.set_use_sharded_ddp(config.get("use_sharded_ddp", False))
)
if not test_only:
task.set_optimizer(optimizer)
task.set_optimizer_schedulers(param_schedulers)
use_gpu = config.get("use_gpu")
if use_gpu is not None:
task.set_use_gpu(use_gpu)
for phase_type in datasets:
task.set_dataset(datasets[phase_type], phase_type)
# NOTE: this is a private member and only meant to be used for
# logging/debugging purposes. See __repr__ implementation
task._config = config
return task
@property
def num_batches_per_phase(self):
"""Returns number of batches in current phase iterator"""
return len(self.data_iterator)
@property
def model(self):
"""Returns model used in training (can be wrapped with DDP)"""
return (
self.distributed_model if is_distributed_training_run() else self.base_model
)
@property
def loss(self):
"""Returns loss used in training (can be wrapped with DDP)"""
return self.distributed_loss if self.distributed_loss else self.base_loss
@property
def phase_type(self):
"""Returns current phase type. String with value "train" or "test" """
return "train" if self.train else "test"
@property
def eval_phase_idx(self):
"""Returns current evaluation phase"""
return self.phase_idx - self.train_phase_idx - 1
def get_total_training_phases(self):
"""
Returns the total number of "train" phases in the task
"""
num_training_phases = 0
for phase in self.phases:
if phase["train"] is True:
num_training_phases += 1
return num_training_phases
def get_total_test_phases(self):
"""
Returns the total number of "test" phases in the task
"""
num_test_phases = 0
for phase in self.phases:
if phase["train"] is False:
num_test_phases += 1
return num_test_phases
def _build_phases(self):
"""Returns list of phases from config.
These phases will look like:
{
train: is this a train or test phase?
optimizer: optimizer settings
}
- If this is a test only run, then only test phases will be
generated
- If this is a training run with both train and test datasets, then x phases =
x train phases + x test phases, interleaved. If test_phase_period > 1, test
phases are only added after test_phase_period train phases. The last phase is
always a test phase.
- If this is a training run with only a train dataset, then x phases = x train
phases.
"""
if not self.test_only:
phases = [
{"train": True}
for _ in range(math.ceil(self.train_phases_per_epoch * self.num_epochs))
]
if self._train_only:
return phases
final_phases = []
for i, phase in enumerate(phases):
final_phases.append(phase)
if (i + 1) % self.test_phase_period == 0:
final_phases.append({"train": False})
if final_phases[-1]["train"]:
final_phases.append({"train": False})
return final_phases
return [{"train": False} for _ in range(self.num_epochs)]
def build_dataloader_from_dataset(self, dataset, **kwargs):
"""Builds a dataloader from the provided dataset
Args:
dataset: A ClassyDataset
kwargs: Additional kwargs to pass during dataloader construction for
derived classes
"""
return dataset.iterator(
phase_type=self.phase_type,
current_phase_id=self.train_phase_idx if self.train else 0,
pin_memory=self.use_gpu and torch.cuda.device_count() > 1,
multiprocessing_context=mp.get_context(self.dataloader_mp_context),
**kwargs,
)
def build_dataloaders_for_current_phase(self):
"""Builds dataloader(s) for the current phase.
Deriving classes can override this method to support custom behavior, like
supporting multiple dataloaders in parallel.
"""
self.dataloader = self.build_dataloader_from_dataset(
self.datasets[self.phase_type]
)
def prepare_optimizer(self, optimizer, model, loss=None):
bn_params, other_params = split_batchnorm_params(model)
if loss is not None:
bn_params_loss, params_loss = split_batchnorm_params(loss)
bn_params = bn_params + bn_params_loss
other_params = other_params + params_loss
bn_schedulers = self.optimizer_schedulers.copy()
if not self.bn_weight_decay:
bn_schedulers["weight_decay"] = 0
param_groups = [{"params": other_params, **self.optimizer_schedulers}]
if len(bn_params) > 0:
param_groups.append({"params": bn_params, **bn_schedulers})
self.optimizer.set_param_groups(param_groups)
def prepare(self):
"""Prepares task for training, populates all derived attributes """
self.phases = self._build_phases()
self.train = False if self.test_only else self.train
if self.batch_norm_sync_mode == BatchNormSyncMode.PYTORCH:
self.base_model = nn.SyncBatchNorm.convert_sync_batchnorm(self.base_model)
elif self.batch_norm_sync_mode == BatchNormSyncMode.APEX:
sync_bn_process_group = apex.parallel.create_syncbn_process_group(
self.batch_norm_sync_group_size
)
self.base_model = apex.parallel.convert_syncbn_model(
self.base_model, process_group=sync_bn_process_group
)
# move the model and loss to the right device
if self.use_gpu:
self.base_model, self.base_loss = copy_model_to_gpu(
self.base_model, self.base_loss
)
else:
self.base_loss.cpu()
self.base_model.cpu()
if self.optimizer is not None:
self.prepare_optimizer(
optimizer=self.optimizer, model=self.base_model, loss=self.base_loss
)
if self.amp_args is not None:
if self.amp_type == AmpType.APEX:
# Initialize apex.amp. This updates the model and the PyTorch optimizer (
# if training, which is wrapped by the ClassyOptimizer in self.optimizer).
# Please note this must happen before loading the checkpoint, cause
# there's amp state to be restored.
if self.optimizer is None:
self.base_model = apex.amp.initialize(
self.base_model, optimizers=None, **self.amp_args
)
else:
self.base_model, self.optimizer.optimizer = apex.amp.initialize(
self.base_model, self.optimizer.optimizer, **self.amp_args
)
if self.simulated_global_batchsize is not None:
if self.simulated_global_batchsize % self.get_global_batchsize() != 0:
raise ValueError(
f"Global batch size ({self.get_global_batchsize()}) must divide "
f"simulated_global_batchsize ({self.simulated_global_batchsize})"
)
else:
self.simulated_global_batchsize = self.get_global_batchsize()
self.optimizer_period = (
self.simulated_global_batchsize // self.get_global_batchsize()
)
if self.optimizer_period > 1:
logging.info(
f"Using gradient accumulation with a period of {self.optimizer_period}"
)
if self.checkpoint_path:
self.checkpoint_dict = load_and_broadcast_checkpoint(self.checkpoint_path)
classy_state_dict = (
None
if self.checkpoint_dict is None
else self.checkpoint_dict["classy_state_dict"]
)
if classy_state_dict is not None:
state_load_success = update_classy_state(self, classy_state_dict)
assert (
state_load_success
), "Update classy state from checkpoint was unsuccessful."
self.init_distributed_data_parallel_model()
def init_distributed_data_parallel_model(self):
"""
Initialize
`torch.nn.parallel.distributed.DistributedDataParallel <https://pytorch.org/
docs/stable/nn.html#distributeddataparallel>`_.
Needed for distributed training. This is where a model should be wrapped by DDP.
"""
if not is_distributed_training_run():
return
assert (
self.distributed_model is None
), "init_ddp_non_elastic must only be called once"
broadcast_buffers = (
self.broadcast_buffers_mode == BroadcastBuffersMode.FORWARD_PASS
)
if self.use_sharded_ddp:
if not isinstance(self.optimizer, ZeRO):
raise ValueError(
"ShardedDataParallel engine should only be used in conjunction with ZeRO optimizer"
)
from fairscale.nn.data_parallel import ShardedDataParallel
# Replace the original DDP wrap by the shard-aware ShardedDDP
self.distributed_model = ShardedDataParallel(
module=self.base_model,
sharded_optimizer=self.optimizer.optimizer,
broadcast_buffers=broadcast_buffers,
)
else:
self.distributed_model = init_distributed_data_parallel_model(
self.base_model,
broadcast_buffers=broadcast_buffers,
find_unused_parameters=self.find_unused_parameters,
bucket_cap_mb=self.ddp_bucket_cap_mb,
)
if self.fp16_grad_compress:
from torch.distributed.algorithms import ddp_comm_hooks
# FP16 hook is stateless and only takes a process group as the state.
# We use the default process group so we set the state to None.
process_group = None
self.distributed_model.register_comm_hook(
process_group,
ddp_comm_hooks.default_hooks.fp16_compress_hook,
)
if (
isinstance(self.base_loss, ClassyLoss)
and self.base_loss.has_learned_parameters()
):
logging.info("Initializing distributed loss")
self.distributed_loss = init_distributed_data_parallel_model(
self.base_loss,
broadcast_buffers=broadcast_buffers,
find_unused_parameters=self.find_unused_parameters,
bucket_cap_mb=self.ddp_bucket_cap_mb,
)
@property
def where(self):
"""Returns the proportion of training that has completed. If in test
only mode, returns proportion of testing completed
Returned value is a float in the range [0, 1)
"""
current_step = self.num_updates / self.get_global_batchsize()
num_phases = (
self.get_total_test_phases()
if self.test_only
else self.get_total_training_phases()
)
if self.num_batches_per_phase <= 0:
raise RuntimeError("No batches to read. Is the dataset empty?")
num_steps = num_phases * self.num_batches_per_phase
where = current_step / num_steps
return where
def get_classy_state(self, deep_copy: bool = False):
"""Returns serialiable state of task
Args:
deep_copy: If true, does a deep copy of state before returning.
"""
optimizer_state = {}
if self.optimizer is not None:
optimizer_state = self.optimizer.get_classy_state()
classy_state_dict = {
"train": self.train,
"base_model": self.base_model.get_classy_state(),
"meters": [meter.get_classy_state() for meter in self.meters],
"optimizer": optimizer_state,
"phase_idx": self.phase_idx,
"train_phase_idx": self.train_phase_idx,
"num_updates": self.num_updates,
"losses": self.losses,
"hooks": {hook.name(): hook.get_classy_state() for hook in self.hooks},
"loss": {},
}
if "train" in self.datasets and self._is_checkpointable_dataset(
self.datasets["train"]
):
classy_state_dict["train_dataset_iterator"] = self.datasets[
"train"
].get_classy_state()
if isinstance(self.base_loss, ClassyLoss):
classy_state_dict["loss"] = self.base_loss.get_classy_state()
if self.amp_args is not None:
if self.amp_type == AmpType.APEX:
classy_state_dict["amp"] = apex.amp.state_dict()
elif self.amp_grad_scaler is not None:
classy_state_dict["amp"] = self.amp_grad_scaler.state_dict()
if deep_copy:
classy_state_dict = copy.deepcopy(classy_state_dict)
return classy_state_dict
def set_classy_state(self, state):
"""Set task state
Args:
state: Dict containing state of a task
"""
# some settings are different in test only
self.train = False if self.test_only else state["train"]
if not self.test_only:
self.phase_idx = state["phase_idx"]
self.num_updates = state["num_updates"]
self.train_phase_idx = state["train_phase_idx"]
self.losses = state["losses"]
for meter, meter_state in zip(self.meters, state["meters"]):
meter.set_classy_state(meter_state)
self.base_model.set_classy_state(state["base_model"])
if self.optimizer is not None:
self.optimizer.set_classy_state(state["optimizer"])
if state.get("loss") and isinstance(self.base_loss, ClassyLoss):
self.base_loss.set_classy_state(state["loss"])
if "amp" in state:
if self.amp_type == AmpType.APEX:
apex.amp.load_state_dict(state["amp"])
else:
self.amp_grad_scaler.load_state_dict(state["amp"])
for hook in self.hooks:
# we still want to be able to run when new hooks are added or old
# hooks are removed
if hook.name() in state["hooks"]:
hook.set_classy_state(state["hooks"][hook.name()])
else:
logging.warning(f"No state found for hook: {hook.name()}")
if "train" in self.datasets and self._is_checkpointable_dataset(
self.datasets["train"]
):
self.datasets["train"].set_classy_state(state.get("train_dataset_iterator"))
@staticmethod
def _is_checkpointable_dataset(dataset):
return hasattr(dataset, "get_classy_state") and hasattr(
dataset, "set_classy_state"
)
def eval_step(self):
self.last_batch = None
# Process next sample
with Timer() as timer:
sample = next(self.data_iterator)
assert isinstance(sample, dict) and "input" in sample and "target" in sample, (
f"Returned sample [{sample}] is not a map with 'input' and"
+ "'target' keys"
)
target = sample["target"]
if self.use_gpu:
sample = recursive_copy_to_gpu(sample, non_blocking=True)
# Optional Pytorch AMP context
torch_amp_context = (
torch.cuda.amp.autocast()
if self.amp_type == AmpType.PYTORCH
else contextlib.suppress()
)
with torch.no_grad(), torch_amp_context:
output = self.model(sample["input"])
local_loss = self.compute_loss(output, sample)
loss = local_loss.detach().clone()
self.check_inf_nan(loss)
self.losses.append(loss.data.cpu().item())
self.update_meters(output, sample)
# Move some data to the task so hooks get a chance to access it
self.last_batch = LastBatchInfo(
loss=loss,
output=output,
target=target,
sample=sample,
step_data={"sample_fetch_time": timer.elapsed_time},
)
def check_inf_nan(self, loss):
if loss == float("inf") or loss == float("-inf") or loss != loss:
raise FloatingPointError(f"Loss is infinity or NaN: {loss}")
def _should_do_step(self):
"""Tells if we will be performing an optimizer step.
Returns True always if there is no gradient accumulation. With gradient
accumulation returns True only when the gradients will be synchronized and we
will be performing an optimizer step.
"""
update_idx = self.num_updates // self.get_global_batchsize()
return (update_idx % self.optimizer_period) == self.optimizer_period - 1
def train_step(self):
"""Train step to be executed in train loop."""
self.last_batch = None
# Process next sample
with Timer() as timer:
sample = next(self.data_iterator)
assert isinstance(sample, dict) and "input" in sample and "target" in sample, (
f"Returned sample [{sample}] is not a map with 'input' and"
+ "'target' keys"
)
# Copy sample to GPU
target = sample["target"]
if self.use_gpu:
sample = recursive_copy_to_gpu(sample, non_blocking=True)
if self.mixup_transform is not None:
sample = self.mixup_transform(sample)
# Optional Pytorch AMP context
torch_amp_context = (
torch.cuda.amp.autocast()
if self.amp_type == AmpType.PYTORCH
else contextlib.suppress()
)
# only sync with DDP when we need to perform an optimizer step
# an optimizer step can be skipped if gradient accumulation is enabled
do_step = self._should_do_step()
ctx_mgr_model = (
self.distributed_model.no_sync()
if self.distributed_model is not None and not do_step
else contextlib.suppress()
)
ctx_mgr_loss = (
self.distributed_loss.no_sync()
if self.distributed_loss is not None and not do_step
else contextlib.suppress()
)
with ctx_mgr_model, ctx_mgr_loss:
# Forward pass
with torch.enable_grad(), torch_amp_context:
output = self.model(sample["input"])
local_loss = self.compute_loss(output, sample)
loss = local_loss.detach().clone()
self.losses.append(loss.data.cpu().item())
self.update_meters(output, sample)
# Backwards pass + optimizer step
self.run_optimizer(local_loss)
self.num_updates += self.get_global_batchsize()
# Move some data to the task so hooks get a chance to access it
self.last_batch = LastBatchInfo(
loss=loss,
output=output,
target=target,
sample=sample,
step_data={"sample_fetch_time": timer.elapsed_time},
)
def compute_loss(self, model_output, sample):
return self.loss(model_output, sample["target"])
def run_optimizer(self, loss):
"""Runs backwards pass and update the optimizer"""
self.check_inf_nan(loss)
# Gradient accumulation logic. We always set optimizer_period, even
# if gradient accumulation is disabled. Assumes all batches have the
# same size
update_idx = self.num_updates // self.get_global_batchsize()
do_zero_grad = (update_idx % self.optimizer_period) == 0
do_step = self._should_do_step()
if do_zero_grad:
self.optimizer.zero_grad()
if self.amp_type == AmpType.APEX:
with apex.amp.scale_loss(loss, self.optimizer.optimizer) as scaled_loss:
scaled_loss.backward()
elif self.amp_type == AmpType.PYTORCH:
self.amp_grad_scaler.scale(loss).backward()
else:
loss.backward()
if do_step:
# Handle gradient accumulation related gradient rescaling
if self.optimizer_period != 1:
self._rescale_gradients(1 / self.optimizer_period)
# Clipping must happen after grad accumulation
if self.clip_grad_norm is not None:
self._clip_gradients(self.clip_grad_norm)
if self.amp_type == AmpType.PYTORCH:
# If using mixed precision, handle underflow-related scaling
# See https://pytorch.org/docs/stable/amp.html#gradient-scaling
# for context
self.amp_grad_scaler.step(self.optimizer, where=self.where)
self.amp_grad_scaler.update()
else:
self.optimizer.step(where=self.where)
def _rescale_gradients(self, scale):
for param in master_params(self.optimizer):
if param.grad is not None:
param.grad.data.mul_(scale)
def _clip_gradients(self, max_norm):
nn.utils.clip_grad_norm_(master_params(self.optimizer), max_norm)
def update_meters(self, model_output, sample):
target = sample["target"].detach().cpu()
model_output = model_output.detach().cpu()
# Update meters
for meter in self.meters:
meter.update(model_output, target, is_train=self.train)
def synchronize_losses(self):
"""Average the losses across the different replicas"""
# Average losses across nodes
losses_tensor = torch.tensor(self.losses)
synchronized_losses_tensor = all_reduce_mean(losses_tensor)
self.losses = synchronized_losses_tensor.tolist()
def advance_phase(self):
"""Performs bookkeeping / task updates between phases
Increments phase idx, resets meters, resets loss history,
resets counters, shuffles dataset, rebuilds iterators, and
sets the train / test state for phase.
"""
logging.debug("Advancing phase")
# Reset meters for next phase / epoch
for meter in self.meters:
meter.reset()
# Reset loss history for next epoch
self.losses = []
# Setup new phase
self.phase_idx += 1
phase = self.phases[self.phase_idx]
self.train = True if phase["train"] else False
if self.train:
self.train_phase_idx += 1
# Re-build dataloader & re-create iterator anytime membership changes.
self.build_dataloaders_for_current_phase()
self.create_data_iterators()
# Set up pytorch module in train vs eval mode, update optimizer.
self._set_model_train_mode()
def done_training(self):
"""Stop condition for training"""
return self.phase_idx + 1 >= len(self.phases)
def create_data_iterators(self):
"""Creates data iterator(s) for the current phase."""
# Delete iterator explicitly so that all dataloader processes
# are cleaned up.
del self.data_iterator
self.data_iterator = iter(self.dataloader)
def _set_model_train_mode(self):
"""Set train mode for model"""
phase = self.phases[self.phase_idx]
self.base_model.train(phase["train"])
self.base_loss.train(phase["train"])
if (
self.broadcast_buffers_mode == BroadcastBuffersMode.BEFORE_EVAL
and not self.train
):
self._broadcast_buffers()
def _broadcast_buffers(self):
"""Explicitly synchronize buffers across all devices."""
if self.distributed_model is None:
return
buffers = list(self.base_model.buffers())
if len(buffers) > 0:
logging.info("Synchronizing buffers before evaluation.")
for buffer in buffers:
broadcast(buffer, 0, group=self.distributed_model.process_group)
# TODO: Functions below should be better abstracted into the dataloader
# abstraction
def get_batchsize_per_replica(self):
"""Return local replica's batchsize for dataset (e.g. batchsize per GPU)"""
return self.datasets[self.phase_type].get_batchsize_per_replica()
def get_global_batchsize(self):
"""Return global batchsize across all trainers"""
return self.datasets[self.phase_type].get_global_batchsize()
def on_start(self):
for hook in self.hooks:
hook.on_start(self)
def on_phase_start(self):
self.phase_start_time_total = time.perf_counter()
self.advance_phase()
for hook in self.hooks:
hook.on_phase_start(self)
self.phase_start_time_train = time.perf_counter()
def on_phase_end(self):
self.log_phase_end("train")
if self.train:
self.optimizer.on_epoch(where=self.where)
logging.debug("Syncing losses on phase end...")
self.synchronize_losses()
logging.debug("...losses synced")
logging.debug("Syncing meters on phase end...")
for meter in self.meters:
meter.sync_state()
logging.debug("...meters synced")
barrier()
for hook in self.hooks:
hook.on_phase_end(self)
self.perf_log = []
self.log_phase_end("total")
def on_end(self):
for hook in self.hooks:
hook.on_end(self)
def log_phase_end(self, tag):
if not self.train:
return
start_time = (
self.phase_start_time_train
if tag == "train"
else self.phase_start_time_total
)
phase_duration = time.perf_counter() - start_time
im_per_sec = (
self.get_global_batchsize() * self.num_batches_per_phase
) / phase_duration
self.perf_log.append(
{
"tag": tag,
"phase_idx": self.train_phase_idx,
"epoch_duration": phase_duration,
"im_per_sec": im_per_sec,
}
)
def __repr__(self):
if hasattr(self, "_config"):
config = json.dumps(self._config, indent=4)
return f"{super().__repr__()} initialized with config:\n{config}"
return super().__repr__()
# create DataLoaders with samplers
train_dataloader = DataLoader(train_ds,
batch_size=100,
sampler=train_sampler,
shuffle=False)
valid_dataloader = DataLoader(valid_ds, batch_size=100, shuffle=True)
test_dataloader = DataLoader(test_ds, batch_size=100, shuffle=True)
# set LR scheduler
scheduler = torch.optim.lr_scheduler.OneCycleLR(
optimizer, max_lr=0.01, total_steps=len(train_dataloader) * epochs)
# create gradient scaler for mixed precision
scaler = GradScaler()
# train function
def train(dataloader):
pbar = ProgressBar(n_total=len(dataloader), desc='Training')
train_loss = AverageMeter()
model.train()
for batch_idx, batch in enumerate(dataloader):
b_features, b_target, b_idx = batch['features'].to(
DEVICE), batch['target'].to(DEVICE), batch['idx'].to(DEVICE)
optimizer.zero_grad()
with autocast():
logits, probs = model(b_features)
loss = F.cross_entropy(logits, b_target)
scaler.scale(loss).backward()
def set_amp_args(self, amp_args: Optional[Dict[str, Any]]):
"""Disable / enable apex.amp and set the automatic mixed precision parameters.
apex.amp can be utilized for mixed / half precision training.
Args:
amp_args: Dictionary containing arguments to be passed to
amp.initialize. Set to None to disable amp. To enable mixed
precision training, pass amp_args={"opt_level": "O1"} here.
See https://nvidia.github.io/apex/amp.html for more info.
Raises:
RuntimeError: If opt_level is not None and apex is not installed.
Warning: apex needs to be installed to utilize this feature.
"""
self.amp_args = amp_args
if amp_args is None:
logging.info("AMP disabled")
else:
# Check that the requested AMP type is known
try:
self.amp_type = AmpType[self.amp_args["amp_type"].upper()]
except KeyError:
logging.info("AMP type not specified, defaulting to Apex")
self.amp_type = AmpType.APEX
# Check for CUDA availability, required for both Apex and Pytorch AMP
if not torch.cuda.is_available():
raise RuntimeError(
"AMP is required but CUDA is not supported, cannot enable AMP"
)
# Check for Apex availability
if self.amp_type == AmpType.APEX and not apex_available:
raise RuntimeError(
"Apex AMP is required but Apex is not installed, cannot enable AMP"
)
if self.use_sharded_ddp:
if self.amp_type == AmpType.APEX:
raise RuntimeError(
"ShardedDDP has been requested, which is incompatible with Apex AMP"
)
if not fairscale_available:
raise RuntimeError(
"ShardedDDP has been requested, but fairscale is not installed in the current environment"
)
# Set Torch AMP grad scaler, used to prevent gradient underflow
elif self.amp_type == AmpType.PYTORCH:
if self.use_sharded_ddp:
logging.info("Using ShardedGradScaler to manage Pytorch AMP")
self.amp_grad_scaler = ShardedGradScaler()
else:
self.amp_grad_scaler = TorchGradScaler()
logging.info(f"AMP enabled with args {amp_args}")
return self
class Learner(object):
def __init__(self,
model,
optimizer,
loss_func,
name="",
scheduler=None,
device='cpu'):
self.model = model
self.optimizer = optimizer
self.loss_func = loss_func
self.scheduler = scheduler
self.scaler = None
self.device = device
self.metric = None
self.name = name
self.log = {}
self.eth = 0.99
self.do_autocast = False
def init_amp(self,
init_scale=65536.0,
growth_factor=2.0,
backoff_factor=0.5,
growth_interval=2000,
enabled=True,
do_autocast=True):
self.do_autocast = do_autocast
if GradScaler is not None:
self.scaler = GradScaler(init_scale=init_scale,
growth_factor=growth_factor,
backoff_factor=backoff_factor,
growth_interval=growth_interval,
enabled=True)
def get_y(self, batch):
# get Y from Batch, the default is batch[-1] but you can overwrite it
return batch[-1]
def get_inds(self, batch):
# get Y from Batch, the default is batch[-1] but you can overwrite it
return batch[-1]
def get_x(self, batch):
# get x from Batch, the default is batch[:-1] but you can overwrite it
if isinstance(batch, (list, tuple)):
return batch[:-1]
else:
return [batch]
def run_model(self, model, batch):
return model(*(x.to(self.device) for x in self.get_x(batch)))
def calc_loss(self, y_pred, y_true):
return self.loss_func(y_pred, y_true.to(self.device))
def one_cycle(self, batch, train=True, do_step=True):
device = self.device
self.preprocess_batch(batch, train)
y_true = self.get_y(batch)
if autocast is None:
y_pred = self.run_model(self.model, batch)
loss = self.calc_loss(y_pred, y_true)
loss_item = 0 if np.isnan(loss.item()) else loss.item()
else:
with autocast(self.do_autocast):
y_pred = self.run_model(self.model, batch)
loss = self.calc_loss(y_pred, y_true)
loss_item = 0 if np.isnan(loss.item()) else loss.item()
if train:
if self.scaler is not None:
self.scaler.scale(loss).backward()
else:
loss.backward()
if do_step:
if self.scaler is not None:
self.scaler.step(self.optimizer)
self.scaler.update()
else:
self.optimizer.step()
if self.scheduler is not None:
self.scheduler.step()
self.optimizer.zero_grad()
if np.isnan(loss.item()):
print('got loss = nan')
loss_item = 0 if np.isnan(loss.item()) else loss.item()
return loss_item if train else (loss_item, y_pred.to('cpu').detach())
def one_training_epoch(self, dl, accumulation_steps=1):
device = self.device
torch.cuda.empty_cache()
avg_loss = 0.
lossf = 0.
self.model = self.model.train()
self.model.zero_grad()
tk0 = notebook.tqdm(dl)
for i, batch in enumerate(tk0):
do_step = (i + 1) % accumulation_steps == 0
loss_item = self.one_cycle(batch, train=True, do_step=do_step)
e = min(self.eth, 1 - 1.0 / (i + 1.0))
lossf = e * lossf + (1 - e) * loss_item
tk0.set_postfix(loss=lossf)
avg_loss += loss_item / len(dl)
tk0.disable = False
tk0.set_postfix(loss=avg_loss)
tk0.disable = True
return avg_loss
def agg_tta(self, y):
return np.stack(y,0).mean(0) if not isinstance(y[0],tuple)\
else tuple(np.stack([yy[i] for yy in y],0).mean(0) for i in range(len(y[0])))
def preprocess_batch(self, batch, train=True):
return (batch)
def one_eval_epoch(self, dl, tta=1):
device = self.device
avg_loss = 0.
avg_accuracy = 0.
lossf = 0
self.model = self.model.eval()
predss = []
with torch.no_grad():
for t in range(tta):
pred_list = []
true_list = []
tk0 = notebook.tqdm(dl)
for i, batch in enumerate(tk0):
loss_item, y_pred = self.one_cycle(batch,
train=False,
do_step=False)
pred_list.append(y_pred.to('cpu').numpy() if not isinstance(y_pred,tuple) else\
tuple(y.to('cpu').numpy() for y in y_pred))
y_batch = self.get_y(batch)
true_list.append(y_batch.to('cpu').numpy() if not isinstance(y_batch,tuple) else\
tuple(y.to('cpu').numpy() for y in y_batch))
e = min(self.eth, 1 - 1.0 / (i + 1.0))
lossf = e * lossf + (1 - e) * loss_item
tk0.set_postfix(loss=lossf)
avg_loss += loss_item / len(dl)
# y_true=np.concatenate(true_list,0)
y_true=np.concatenate(true_list,0) if not isinstance(true_list[0],tuple) else\
tuple(np.concatenate([p[i] for p in true_list],0) for i in range(len(true_list[0])))
predss.append(np.concatenate(pred_list,0) if not isinstance(pred_list[0],tuple) else\
tuple(np.concatenate([p[i] for p in pred_list],0) for i in range(len(pred_list[0]))))
preds = self.agg_tta(predss, 0) if tta > 1 else predss[0]
m = dict() if self.metric is None else self.metric(preds, y_true)
tk0.disable = False
tk0.set_postfix(loss=avg_loss, **m)
tk0.disable = True
return avg_loss, m
def send_log(self, **kwargs):
log = {'model': self.name}
log.update(kwargs)
try:
sandesh.send(log)
except:
print(log)
def save2log(self, **kwargs):
for key in kwargs.keys():
if key not in self.log:
self.log[key] = []
self.log[key].append(kwargs[key])
def evaluate(self, ds, num_workers=8, tta=1, dl_args={'shuffle': False}):
dl = D.DataLoader(ds, num_workers=num_workers, **dl_args)
return self.one_eval_epoch(dl, tta=tta)
def fit(self,
num_epoches,
train_ds,
validate_ds=None,
batch_size=None,
lr=None,
accumulation_steps=1,
num_workers=8,
send_log=True,
eval_batch=None,
reset_best=False,
make_best=True,
tta=1,
train_dl_args={'shuffle': True},
val_dl_args={'shuffle': False},
save_checkpoint='best',
path=''):
if batch_size is not None:
train_dl_args['batch_size'] = batch_size
val_dl_args['batch_size'] = batch_size
if eval_batch is not None:
val_dl_args['batch_size'] = eval_batch
tq = notebook.tqdm(range(num_epoches))
if lr is not None:
self.set_lr(lr)
if reset_best or not hasattr(self, 'best_metric'):
self.best_model = None
self.best_metric = np.inf
for k, epoch in enumerate(tq):
self.on_epoch_begin(epoch,
train_ds=train_ds,
validate_ds=validate_ds)
dl = D.DataLoader(train_ds,
num_workers=num_workers,
**train_dl_args)
if next(self.model.parameters()).device != torch.device('cpu'):
torch.cuda.empty_cache()
tavg_loss = self.one_training_epoch(
dl, accumulation_steps=accumulation_steps)
# dl=D.DataLoader(validate_ds, batch_size=batch_size if eval_batch is None else eval_batch,
# num_workers=num_workers,**val_dl_args)
if validate_ds is not None:
avg_loss, metric = self.evaluate(validate_ds,
num_workers=num_workers,
dl_args=val_dl_args,
tta=tta)
else:
avg_loss = tavg_loss
metric = {}
if send_log:
self.send_log(epoch=epoch,
tloss=tavg_loss,
loss=avg_loss,
**metric)
self.save2log(epoch=epoch,
tloss=tavg_loss,
loss=avg_loss,
**metric)
m = avg_loss if 'metric' not in metric.keys() else metric['metric']
if save_checkpoint == 'last':
self.save_checkpoint(path)
if m < self.best_metric:
self.best_metric = m
self.best_model = copy.deepcopy(self.model.state_dict())
tq.set_postfix(best_metric=self.best_metric)
if save_checkpoint == 'best':
self.save_checkpoint(path)
self.on_epoch_end(epoch)
print('best metric:', self.best_metric)
if make_best:
self.model.load_state_dict(self.best_model)
def save_model(self, path, name=None):
name = self.name if name is None else name
torch.save(self.model.state_dict(), f'{path}{name}')
def load_model(self, path, name=None, map_location=None):
name = self.name if name is None else name
self.model.load_state_dict(
torch.load(f'{path}{name}', map_location=map_location))
def save_checkpoint(self, path, name=None):
name = self.name + '.chk' if name is None else name
checkpoint = {
'model': self.model.state_dict(),
'best_model': self.best_model,
'best_metric': self.best_metric,
'model': self.model.state_dict(),
'optimizer': self.optimizer.state_dict(),
'log': self.log
}
if self.scaler:
checkpoint['scaler'] = self.scaler.state_dict()
torch.save(checkpoint, f'{path}{name}')
def load_checkpoint(self, path, name=None):
name = self.name + '.chk' if name is None else name + '.chk'
checkpoint = torch.load(f'{path}{name}')
self.model.load_state_dict(checkpoint['model'])
self.best_model = checkpoint['best_model']
self.best_metric = checkpoint['best_metric']
self.optimizer.load_state_dict(checkpoint['optimizer'])
self.log = checkpoint['log']
if 'scaler' in checkpoint.keys():
self.scaler = GradScaler()
self.scaler.load_state_dict(checkpoint['scaler'])
else:
self.scaler = None
def set_lr(self, lr):
for param_group in self.optimizer.param_groups:
param_group['lr'] = lr
def on_epoch_begin(self, *args, **kargs):
pass
def on_epoch_end(self, *args, **kargs):
pass
def predict(self,
ds,
batch_size=None,
num_workers=8,
dl_args={'shuffle': False},
return_inds=False,
return_true=False,
verbose=True,
do_eval=True):
device = self.device
if batch_size is not None:
dl_args['batch_size'] = batch_size
dl = D.DataLoader(ds, num_workers=num_workers, **dl_args)
pred_list = []
inds_list = []
true_list = []
if do_eval:
self.model = self.model.eval()
with torch.no_grad():
tk0 = notebook.tqdm(dl) if verbose else dl
for i, batch in enumerate(tk0):
if autocast is None:
y_pred = self.run_model(self.model, batch)
else:
with autocast(self.scaler is not None):
y_pred = self.run_model(self.model, batch)
if return_inds:
inds_list.append(self.get_inds(batch).to('cpu').numpy())
if return_true:
yb = self.get_y(batch)
true_list.append(yb.to('cpu').numpy() if not isinstance(yb,tuple) else\
tuple(y.to('cpu').numpy() for y in yb))
pred_list.append(y_pred.to('cpu').numpy() if not isinstance(y_pred,tuple) else\
tuple(y.to('cpu').numpy() for y in y_pred))
pred = np.concatenate(pred_list,0) if not isinstance(pred_list[0],tuple) else\
tuple(np.concatenate([p[i] for p in pred_list],0) for i in range(len(pred_list[0])))
out = ()
if return_inds:
out = out + (np.concatenate(inds_list, 0), )
if return_true:
rt=np.concatenate(true_list,0) if not isinstance(true_list[0],tuple) else\
tuple(np.concatenate([p[i] for p in true_list],0) for i in range(len(true_list[0])))
out = out + (rt, )
return pred if len(out) == 0 else (pred, ) + out
def train_schedule(writer, loader, val_num_steps, validation_loader, device,
criterion, net, optimizer, lr_scheduler, num_epochs,
is_mixed_precision, num_classes, categories, input_sizes,
selector, classes, encoder_only):
# Poly training schedule
# Validate and find the best snapshot
best_mIoU = 0
net.train()
epoch = 0
running_loss = 0.0
loss_num_steps = int(len(loader) / 10)
if is_mixed_precision:
scaler = GradScaler()
# Training
while epoch < num_epochs:
net.train()
conf_mat = ConfusionMatrix(num_classes)
time_now = time.time()
for i, data in enumerate(loader, 0):
inputs, labels = data
inputs, labels = inputs.to(device), labels.to(device)
optimizer.zero_grad()
with autocast(is_mixed_precision):
outputs = net(inputs)['out']
if encoder_only:
labels = labels.unsqueeze(0)
if labels.dtype not in (torch.float32, torch.float64):
labels = labels.to(torch.float32)
labels = torch.nn.functional.interpolate(
labels, size=input_sizes[1], mode='nearest')
labels = labels.to(torch.int64)
labels = labels.squeeze(0)
else:
outputs = torch.nn.functional.interpolate(
outputs,
size=input_sizes[0],
mode='bilinear',
align_corners=True)
conf_mat.update(labels.flatten(), outputs.argmax(1).flatten())
loss = criterion(outputs, labels)
if is_mixed_precision:
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
else:
loss.backward()
optimizer.step()
lr_scheduler.step()
running_loss += loss.item()
current_step_num = int(epoch * len(loader) + i + 1)
# Record losses
if current_step_num % loss_num_steps == (loss_num_steps - 1):
print('[%d, %d] loss: %.4f' %
(epoch + 1, i + 1, running_loss / loss_num_steps))
writer.add_scalar('training loss',
running_loss / loss_num_steps,
current_step_num)
running_loss = 0.0
# Validate and find the best snapshot
if current_step_num % val_num_steps == (val_num_steps - 1):
test_pixel_accuracy, test_mIoU = test_one_set(
loader=validation_loader,
device=device,
net=net,
num_classes=num_classes,
categories=categories,
output_size=input_sizes[2],
labels_size=input_sizes[1],
selector=selector,
classes=classes,
is_mixed_precision=is_mixed_precision,
encoder_only=encoder_only)
writer.add_scalar('test pixel accuracy', test_pixel_accuracy,
current_step_num)
writer.add_scalar('test mIoU', test_mIoU, current_step_num)
net.train()
# Record best model (straight to disk)
if test_mIoU > best_mIoU:
best_mIoU = test_mIoU
save_checkpoint(net=net,
optimizer=optimizer,
lr_scheduler=lr_scheduler)
# Evaluate training accuracies (same metric as validation, but must be on-the-fly to save time)
with autocast(is_mixed_precision):
acc_global, acc, iu = conf_mat.compute()
print(categories)
print(('global correct: {:.2f}\n'
'average row correct: {}\n'
'IoU: {}\n'
'mean IoU: {:.2f}').format(
acc_global.item() * 100,
['{:.2f}'.format(i) for i in (acc * 100).tolist()],
['{:.2f}'.format(i) for i in (iu * 100).tolist()],
iu.mean().item() * 100))
train_pixel_acc = acc_global.item() * 100
train_mIoU = iu.mean().item() * 100
writer.add_scalar('train pixel accuracy', train_pixel_acc, epoch + 1)
writer.add_scalar('train mIoU', train_mIoU, epoch + 1)
epoch += 1
print('Epoch time: %.2fs' % (time.time() - time_now))
