class BaseTrainer:
def __init__(self, options):
self.options = options
self.device = torch.device(
'cuda:0' if torch.cuda.is_available() else 'cpu')
# override this function to define your model, optimizers etc.
self._init_fn()
self.saver = CheckpointSaver(save_dir=options.checkpoint_dir)
self.summary_writer = SummaryWriter(self.options.summary_dir)
self.checkpoint = None
if self.options.resume and self.saver.exists_checkpoint():
self.checkpoint = self.saver.load_checkpoint(
self.models_dict,
self.optimizers_dict,
checkpoint_file=self.options.checkpoint)
if self.checkpoint is None:
self.epoch_count = 0
self.step_count = 0
else:
self.epoch_count = self.checkpoint['epoch']
self.step_count = self.checkpoint['total_step_count']
# self.lr_schedulers = {k: torch.optim.lr_scheduler.ReduceLROnPlateau(v, patience=5)
# for k,v in self.optimizers_dict.items()}
self.lr_schedulers = {k: torch.optim.lr_scheduler.ExponentialLR(v, gamma=self.options.lr_decay, last_epoch=self.epoch_count-1)\
for k,v in self.optimizers_dict.items()}
for opt in self.optimizers_dict:
self.lr_schedulers[opt].step()
def _init_fn(self):
raise NotImplementedError('You need to provide an _init_fn method')
# @profile
def train(self):
self.endtime = time.time() + self.options.time_to_run
for epoch in tqdm(range(self.epoch_count, self.options.num_epochs),
total=self.options.num_epochs,
initial=self.epoch_count):
train_data_loader = CheckpointDataLoader(
self.train_ds,
checkpoint=self.checkpoint,
batch_size=self.options.batch_size,
num_workers=self.options.num_workers,
pin_memory=self.options.pin_memory,
shuffle=self.options.shuffle_train)
for step, batch in enumerate(
tqdm(train_data_loader,
desc='Epoch ' + str(epoch),
total=math.ceil(
len(self.train_ds) / self.options.batch_size),
initial=train_data_loader.checkpoint_batch_idx),
train_data_loader.checkpoint_batch_idx):
#if epoch == 1: #step == 74 or step == 73:
#from IPython.core.debugger import Pdb
#Pdb().set_trace()
# print("Epoch", epoch, "Step", step)
# print(batch['keypoint_locs'])
if time.time() < self.endtime:
batch = {k: v.to(self.device) for k, v in batch.items()}
out = self._train_step(batch)
self.step_count += 1
if self.step_count % self.options.summary_steps == 0:
try:
self._train_summaries(batch, *out)
except:
from IPython.core.debugger import Pdb
Pdb().set_trace()
if self.step_count % self.options.checkpoint_steps == 0:
self.saver.save_checkpoint(
self.models_dict, self.optimizers_dict, epoch,
step + 1, self.options.batch_size,
train_data_loader.sampler.dataset_perm,
self.step_count)
tqdm.write('Checkpoint saved')
if self.step_count % self.options.test_steps == 0:
val_loss = self.test()
# for opt in self.optimizers_dict:
# self.lr_schedulers[opt].step(val_loss)
else:
tqdm.write('Timeout reached')
self.saver.save_checkpoint(
self.models_dict, self.optimizers_dict, epoch, step,
self.options.batch_size,
train_data_loader.sampler.dataset_perm,
self.step_count)
tqdm.write('Checkpoint saved')
sys.exit(0)
# apply the learning rate scheduling policy
for opt in self.optimizers_dict:
self.lr_schedulers[opt].step()
# load a checkpoint only on startup, for the next epochs
# just iterate over the dataset as usual
self.checkpoint = None
# save checkpoint after each epoch
if (epoch + 1) % 10 == 0:
# self.saver.save_checkpoint(self.models_dict, self.optimizers_dict, epoch+1, 0, self.step_count)
self.saver.save_checkpoint(self.models_dict,
self.optimizers_dict, epoch + 1, 0,
self.options.batch_size, None,
self.step_count)
return
def _get_lr(self):
return next(iter(self.optimizers_dict.values())).param_groups[0]['lr']
# return next(iter(self.lr_schedulers.values())).get_lr()[0]
def _train_step(self, input_batch):
raise NotImplementedError('You need to provide a _train_step method')
def _train_summaries(self, input_batch):
raise NotImplementedError(
'You need to provide a _save_summaries method')
def test(self, input_batch):
raise NotImplementedError('You need to provide a _test_step method')
class BaseTrainer(object):
"""Base class for Trainer objects.
Takes care of checkpointing/logging/resuming training.
"""
def __init__(self, options):
self.options = options
self.endtime = time.time() + self.options.time_to_run
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# override this function to define your model, optimizers etc.
self.init_fn()
self.saver = CheckpointSaver(save_dir=options.checkpoint_dir)
self.checkpoint = None
if self.options.resume and self.saver.exists_checkpoint():
self.checkpoint = self.saver.load_checkpoint(self.models_dict, self.optimizers_dict, checkpoint_file=self.options.checkpoint)
if self.checkpoint is None:
self.epoch_count = 0
self.step_count = 0
else:
self.epoch_count = self.checkpoint['epoch']
self.step_count = self.checkpoint['total_step_count']
def load_pretrained(self, checkpoint_file=None):
"""Load a pretrained checkpoint.
This is different from resuming training using --resume.
"""
if checkpoint_file is not None:
checkpoint = torch.load(checkpoint_file)
for model in self.models_dict:
if model in checkpoint:
state_dict = checkpoint[model]
renamed_state_dict = OrderedDict()
# change the names in the state_dict to match the new layer
for key, value in state_dict.items():
if 'layer' in key:
names = key.split('.')
names[1:1] = ['hmr_layer']
new_key = '.'.join(n for n in names)
renamed_state_dict[new_key] = value
else:
renamed_state_dict[key] = value
self.models_dict[model].load_state_dict(renamed_state_dict, strict=False)
@staticmethod
def linear_rampup(current, rampup_length):
"""Linear rampup"""
assert current >= 0 and rampup_length >= 0
if current >= rampup_length:
return 1.0
else:
return current / rampup_length
def train(self):
"""Training process."""
ramp_step = 0
# Run training for num_epochs epochs
for epoch in tqdm(range(self.epoch_count, self.options.num_epochs), total=self.options.num_epochs, initial=self.epoch_count):
# ------------------ update image size intervals ----------------------
self.train_ds.update_size_intervals(epoch)
# ---------------------------------------------------------------------
# ------------------ update batch size ----------------------
if epoch == 0:
batch_size = self.options.batch_size # 24
elif epoch == 1:
batch_size = self.options.batch_size // 2 # 12
else:
batch_size = self.options.batch_size // 3 # 8
if epoch == 3:
self.options.checkpoint_steps = 2000
# ---------------------------------------------------------------------
# Create new DataLoader every epoch and (possibly) resume from an arbitrary step inside an epoch
train_data_loader = CheckpointDataLoader(self.train_ds, checkpoint=self.checkpoint,
batch_size=batch_size,
num_workers=self.options.num_workers,
pin_memory=self.options.pin_memory,
shuffle=self.options.shuffle_train)
# init alphas
if epoch <= 3:
self.model.init_alphas(epoch+1, self.device)
# Iterate over all batches in an epoch
for step, batch in enumerate(tqdm(train_data_loader, desc='Epoch '+str(epoch),
total=len(self.train_ds) // batch_size,
initial=train_data_loader.checkpoint_batch_idx),
train_data_loader.checkpoint_batch_idx):
# ------------------ ramp consistency loss weight after updating the scale interval ----------------------
if self.options.ramp == 'up':
total_ramp = (len(self.train_ds) // self.options.batch_size) * 5
self.consistency_loss_ramp = self.linear_rampup(ramp_step, total_ramp)
ramp_step += 1
elif self.options.ramp == 'down':
total_ramp = (len(self.train_ds) // self.options.batch_size) * 5
consistency_loss_ramp = self.linear_rampup(ramp_step, total_ramp)
self.consistency_loss_ramp = 1.0 - consistency_loss_ramp
ramp_step += 1
else:
self.consistency_loss_ramp = 1.0
# ---------------------------------------------------------------------
if time.time() < self.endtime:
batch = {k: v.to(self.device) if isinstance(v, torch.Tensor) and k != 'sample_index' else v for k,v in batch.items()}
out = self.train_step(batch)
self.step_count += 1
# Save checkpoint every checkpoint_steps steps
if self.step_count % self.options.checkpoint_steps == 0 and epoch >= 3:
self.saver.save_checkpoint(self.models_dict, self.optimizers_dict, epoch, step+1, self.options.batch_size, train_data_loader.sampler.dataset_perm, self.step_count)
tqdm.write('Checkpoint saved')
else:
tqdm.write('Timeout reached')
self.saver.save_checkpoint(self.models_dict, self.optimizers_dict, epoch, step, self.options.batch_size, train_data_loader.sampler.dataset_perm, self.step_count)
tqdm.write('Checkpoint saved')
sys.exit(0)
# for the first 3 epochs, we only train half epoch
if epoch == 0:
if (step + 1) == (len(self.train_ds) // (self.options.batch_size * 2)):
break
elif epoch == 1:
if (step + 1) == (len(self.train_ds) // self.options.batch_size):
break
elif epoch == 2:
if (step + 1) == (len(self.train_ds) // (self.options.batch_size * 2)) * 3:
break
# load a checkpoint only on startup, for the next epochs
# just iterate over the dataset as usual
self.checkpoint=None
# update learning rate if lr scheduler is epoch-based
if self.lr_scheduler is not None and isinstance(self.lr_scheduler, torch.optim.lr_scheduler.ExponentialLR):
if (epoch + 1) % 4 == 0:
self.lr_scheduler.step()
return
# The following methods (with the possible exception of test) have to be implemented in the derived classes
def init_fn(self):
raise NotImplementedError('You need to provide an _init_fn method')
def train_step(self, input_batch):
raise NotImplementedError('You need to provide a train_step method')
def train_summaries(self, input_batch):
raise NotImplementedError('You need to provide a _train_summaries method')
def test(self):
pass
def main():
args = parser.parse_args()
if args.output:
output_base = args.output
else:
output_base = './output'
exp_name = '-'.join([
datetime.now().strftime("%Y%m%d-%H%M%S"),
args.model,
args.gp,
'f'+str(args.fold)])
output_dir = get_outdir(output_base, 'train', exp_name)
train_input_root = os.path.join(args.data)
batch_size = args.batch_size
num_epochs = args.epochs
wav_size = (16000,)
num_classes = len(dataset.get_labels())
torch.manual_seed(args.seed)
model = model_factory.create_model(
args.model,
in_chs=1,
pretrained=args.pretrained,
num_classes=num_classes,
drop_rate=args.drop,
global_pool=args.gp,
checkpoint_path=args.initial_checkpoint)
#model.reset_classifier(num_classes=num_classes)
dataset_train = dataset.CommandsDataset(
root=train_input_root,
mode='train',
fold=args.fold,
wav_size=wav_size,
format='spectrogram',
)
loader_train = data.DataLoader(
dataset_train,
batch_size=batch_size,
pin_memory=True,
shuffle=True,
num_workers=args.workers
)
dataset_eval = dataset.CommandsDataset(
root=train_input_root,
mode='validate',
fold=args.fold,
wav_size=wav_size,
format='spectrogram',
)
loader_eval = data.DataLoader(
dataset_eval,
batch_size=args.batch_size,
pin_memory=True,
shuffle=False,
num_workers=args.workers
)
train_loss_fn = validate_loss_fn = torch.nn.CrossEntropyLoss()
train_loss_fn = train_loss_fn.cuda()
validate_loss_fn = validate_loss_fn.cuda()
opt_params = list(model.parameters())
if args.opt.lower() == 'sgd':
optimizer = optim.SGD(
opt_params, lr=args.lr,
momentum=args.momentum, weight_decay=args.weight_decay, nesterov=True)
elif args.opt.lower() == 'adam':
optimizer = optim.Adam(
opt_params, lr=args.lr, weight_decay=args.weight_decay, eps=args.opt_eps)
elif args.opt.lower() == 'nadam':
optimizer = nadam.Nadam(
opt_params, lr=args.lr, weight_decay=args.weight_decay, eps=args.opt_eps)
elif args.opt.lower() == 'adadelta':
optimizer = optim.Adadelta(
opt_params, lr=args.lr, weight_decay=args.weight_decay, eps=args.opt_eps)
elif args.opt.lower() == 'rmsprop':
optimizer = optim.RMSprop(
opt_params, lr=args.lr, alpha=0.9, eps=args.opt_eps,
momentum=args.momentum, weight_decay=args.weight_decay)
else:
assert False and "Invalid optimizer"
del opt_params
if not args.decay_epochs:
print('No decay epoch set, using plateau scheduler.')
lr_scheduler = ReduceLROnPlateau(optimizer, patience=10)
else:
lr_scheduler = None
# optionally resume from a checkpoint
start_epoch = 0 if args.start_epoch is None else args.start_epoch
if args.resume:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
if isinstance(checkpoint, dict) and 'state_dict' in checkpoint:
if 'args' in checkpoint:
print(checkpoint['args'])
new_state_dict = OrderedDict()
for k, v in checkpoint['state_dict'].items():
if k.startswith('module'):
name = k[7:] # remove `module.`
else:
name = k
new_state_dict[name] = v
model.load_state_dict(new_state_dict)
if 'optimizer' in checkpoint:
optimizer.load_state_dict(checkpoint['optimizer'])
if 'loss' in checkpoint:
train_loss_fn.load_state_dict(checkpoint['loss'])
print("=> loaded checkpoint '{}' (epoch {})".format(args.resume, checkpoint['epoch']))
start_epoch = checkpoint['epoch'] if args.start_epoch is None else args.start_epoch
else:
model.load_state_dict(checkpoint)
else:
print("=> no checkpoint found at '{}'".format(args.resume))
exit(1)
saver = CheckpointSaver(checkpoint_dir=output_dir)
if args.num_gpu > 1:
model = torch.nn.DataParallel(model, device_ids=list(range(args.num_gpu))).cuda()
else:
model.cuda()
# Optional fine-tune of only the final classifier weights for specified number of epochs (or part of)
if not args.resume and args.ft_epochs > 0.:
if isinstance(model, torch.nn.DataParallel):
classifier_params = model.module.get_classifier().parameters()
else:
classifier_params = model.get_classifier().parameters()
if args.opt.lower() == 'adam':
finetune_optimizer = optim.Adam(
classifier_params,
lr=args.ft_lr, weight_decay=args.weight_decay)
else:
finetune_optimizer = optim.SGD(
classifier_params,
lr=args.ft_lr, momentum=args.momentum, weight_decay=args.weight_decay, nesterov=True)
finetune_epochs_int = int(np.ceil(args.ft_epochs))
finetune_final_batches = int(np.ceil((1 - (finetune_epochs_int - args.ft_epochs)) * len(loader_train)))
print(finetune_epochs_int, finetune_final_batches)
for fepoch in range(0, finetune_epochs_int):
if fepoch == finetune_epochs_int - 1 and finetune_final_batches:
batch_limit = finetune_final_batches
else:
batch_limit = 0
train_epoch(
fepoch, model, loader_train, finetune_optimizer, train_loss_fn, args,
output_dir=output_dir, batch_limit=batch_limit)
best_loss = None
try:
for epoch in range(start_epoch, num_epochs):
if args.decay_epochs:
adjust_learning_rate(
optimizer, epoch, initial_lr=args.lr,
decay_rate=args.decay_rate, decay_epochs=args.decay_epochs)
train_metrics = train_epoch(
epoch, model, loader_train, optimizer, train_loss_fn, args,
saver=saver, output_dir=output_dir)
# save a recovery in case validation blows up
saver.save_recovery({
'epoch': epoch + 1,
'arch': args.model,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict(),
'loss': train_loss_fn.state_dict(),
'args': args,
'gp': args.gp,
},
epoch=epoch + 1,
batch_idx=0)
step = epoch * len(loader_train)
eval_metrics = validate(
step, model, loader_eval, validate_loss_fn, args,
output_dir=output_dir)
if lr_scheduler is not None:
lr_scheduler.step(eval_metrics['eval_loss'])
rowd = OrderedDict(epoch=epoch)
rowd.update(train_metrics)
rowd.update(eval_metrics)
with open(os.path.join(output_dir, 'summary.csv'), mode='a') as cf:
dw = csv.DictWriter(cf, fieldnames=rowd.keys())
if best_loss is None: # first iteration (epoch == 1 can't be used)
dw.writeheader()
dw.writerow(rowd)
# save proper checkpoint with eval metric
best_loss = saver.save_checkpoint({
'epoch': epoch + 1,
'arch': args.model,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict(),
'args': args,
'gp': args.gp,
},
epoch=epoch + 1,
metric=eval_metrics['eval_loss'])
except KeyboardInterrupt:
pass
print('*** Best loss: {0} (epoch {1})'.format(best_loss[1], best_loss[0]))
class BaseTrainer(object):
"""Base class for Trainer objects.
Takes care of checkpointing/logging/resuming training.
"""
def __init__(self, options):
self.options = options
self.endtime = time.time() + self.options.time_to_run
self.device = torch.device(
'cuda' if torch.cuda.is_available() else 'cpu')
# override this function to define your model, optimizers etc.
self.init_fn()
self.saver = CheckpointSaver(save_dir=options.checkpoint_dir)
self.summary_writer = SummaryWriter(self.options.summary_dir)
self.checkpoint = None
if self.options.resume and self.saver.exists_checkpoint():
self.checkpoint = self.saver.load_checkpoint(
self.models_dict,
self.optimizers_dict,
checkpoint_file=self.options.checkpoint)
if self.checkpoint is None:
self.epoch_count = 0
self.step_count = 0
else:
self.epoch_count = self.checkpoint['epoch']
self.step_count = self.checkpoint['total_step_count']
def load_pretrained(self, checkpoint_file=None):
"""Load a pretrained checkpoint.
This is different from resuming training using --resume.
"""
if checkpoint_file is not None:
checkpoint = torch.load(checkpoint_file)
for model in self.models_dict:
if model in checkpoint:
self.models_dict[model].load_state_dict(checkpoint[model])
print('Checkpoint loaded')
def train(self):
"""Training process."""
# Run training for num_epochs epochs
for epoch in tqdm(range(self.epoch_count, self.options.num_epochs),
total=self.options.num_epochs,
initial=self.epoch_count):
# Create new DataLoader every epoch and (possibly) resume from an arbitrary step inside an epoch
train_data_loader = CheckpointDataLoader(
self.train_ds,
checkpoint=self.checkpoint,
batch_size=self.options.batch_size,
num_workers=self.options.num_workers,
pin_memory=self.options.pin_memory,
shuffle=self.options.shuffle_train)
# Iterate over all batches in an epoch
for step, batch in enumerate(
tqdm(train_data_loader,
desc='Epoch ' + str(epoch),
total=len(self.train_ds) // self.options.batch_size,
initial=train_data_loader.checkpoint_batch_idx),
train_data_loader.checkpoint_batch_idx):
if time.time() < self.endtime:
batch = {
k:
v.to(self.device) if isinstance(v, torch.Tensor) else v
for k, v in batch.items()
}
out = self.train_step(batch)
self.step_count += 1
# Tensorboard logging every summary_steps steps
if self.step_count % self.options.summary_steps == 0:
self.train_summaries(batch, *out)
# Save checkpoint every checkpoint_steps steps
if self.step_count % self.options.checkpoint_steps == 0:
self.saver.save_checkpoint(
self.models_dict, self.optimizers_dict, epoch,
step + 1, self.options.batch_size,
train_data_loader.sampler.dataset_perm,
self.step_count)
tqdm.write('Checkpoint saved')
# Run validation every test_steps steps
if self.step_count % self.options.test_steps == 0:
self.test()
else:
tqdm.write('Timeout reached')
self.saver.save_checkpoint(
self.models_dict, self.optimizers_dict, epoch, step,
self.options.batch_size,
train_data_loader.sampler.dataset_perm,
self.step_count)
tqdm.write('Checkpoint saved')
sys.exit(0)
# load a checkpoint only on startup, for the next epochs
# just iterate over the dataset as usual
self.checkpoint = None
# save checkpoint after each epoch
if (epoch + 1) % 10 == 0:
# self.saver.save_checkpoint(self.models_dict, self.optimizers_dict, epoch+1, 0, self.step_count)
self.saver.save_checkpoint(self.models_dict,
self.optimizers_dict, epoch + 1, 0,
self.options.batch_size, None,
self.step_count)
return
# The following methods (with the possible exception of test) have to be implemented in the derived classes
def init_fn(self):
raise NotImplementedError('You need to provide an _init_fn method')
def train_step(self, input_batch):
raise NotImplementedError('You need to provide a _train_step method')
def train_summaries(self, input_batch):
raise NotImplementedError(
'You need to provide a _train_summaries method')
def test(self):
pass
class BaseTrainer(object):
"""
Base class for Trainer objects.
Takes care of checkpointing/logging/resuming training.
"""
def __init__(self, options):
self.options = options
self.endtime = time.time() + self.options.time_to_run
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# override this function to define your model, optimizers etc.
self.init_fn()
self.saver = CheckpointSaver(save_dir=options.checkpoint_dir)
self.summary_writer = SummaryWriter(self.options.summary_dir)
self.checkpoint = None
if self.options.resume and self.saver.exists_checkpoint():
self.checkpoint = self.saver.load_checkpoint(self.models_dict, self.optimizers_dict, checkpoint_file=self.options.checkpoint)
if self.checkpoint is None:
self.epoch_count = 0
self.step_count = 0
else:
self.epoch_count = self.checkpoint['epoch']
self.step_count = self.checkpoint['total_step_count']
def load_pretrained(self, checkpoint_file=None):
"""Load a pretrained checkpoint.
This is different from resuming training using --resume.
"""
if checkpoint_file is not None:
checkpoint = torch.load(checkpoint_file)
for model in self.models_dict:
if model in checkpoint:
self.models_dict[model].load_state_dict(checkpoint[model])
print('Checkpoint loaded')
def train(self):
"""Training process."""
# Run training for num_epochs epochs
for epoch in range(self.epoch_count, self.options.num_epochs):
# Create new DataLoader every epoch and (possibly) resume from an arbitrary step inside an epoch
train_data_loader = CheckpointDataLoader(self.train_ds, checkpoint=self.checkpoint,
batch_size=self.options.batch_size,
num_workers=self.options.num_workers,
pin_memory=self.options.pin_memory,
shuffle=self.options.shuffle_train)
# Iterate over all batches in an epoch
for step, batch in enumerate(tqdm(train_data_loader, desc='Epoch ' + str(epoch),
total=len(self.train_ds) // self.options.batch_size,
initial=train_data_loader.checkpoint_batch_idx),
train_data_loader.checkpoint_batch_idx):
if time.time() < self.endtime:
batch = {k: v.to(self.device) if isinstance(v, torch.Tensor) else v for k,v in batch.items()}
out = self.train_step(batch)
self.step_count += 1
# Tensorboard logging every summary_steps steps
if self.step_count % self.options.summary_steps == 0:
self.train_summaries(batch, *out)
# Save checkpoint every checkpoint_steps steps
if self.step_count % self.options.checkpoint_steps == 0:
self.saver.save_checkpoint(self.models_dict, self.optimizers_dict, epoch, step+1, self.options.batch_size, train_data_loader.sampler.dataset_perm, self.step_count)
tqdm.write('Checkpoint saved')
# Run validation every test_steps steps
if self.step_count % self.options.test_steps == 0:
self.test()
else:
tqdm.write('Timeout reached')
self.saver.save_checkpoint(self.models_dict, self.optimizers_dict, epoch, step, self.options.batch_size, train_data_loader.sampler.dataset_perm, self.step_count)
tqdm.write('Checkpoint saved')
sys.exit(0)
# load a checkpoint only on startup, for the next epochs
# just iterate over the dataset as usual
self.checkpoint=None
# save checkpoint after each epoch
if (epoch+1) % 10 == 0:
# self.saver.save_checkpoint(self.models_dict, self.optimizers_dict, epoch+1, 0, self.step_count)
self.saver.save_checkpoint(self.models_dict, self.optimizers_dict, epoch+1, 0, self.options.batch_size, None, self.step_count)
return
# The following methods (with the possible exception of test) have to be implemented in the derived classes
def init_fn(self):
raise NotImplementedError('You need to provide an _init_fn method')
def train_step(self, input_batch):
raise NotImplementedError('You need to provide a _train_step method')
def train_summaries(self, input_batch):
raise NotImplementedError('You need to provide a _train_summaries method')
def test(self):
pass
def error_adaptive_weight(self, fit_joint_error):
weight = (1 - 10 * fit_joint_error)
weight[weight <= 0] = 0
return weight
def keypoint_loss(self, pred_keypoints_2d, gt_keypoints_2d, weight=None):
"""
Compute 2D reprojection loss on the keypoints.
The loss is weighted by the weight
The available keypoints are different for each dataset.
"""
if gt_keypoints_2d.shape[2] == 3:
conf = gt_keypoints_2d[:, :, -1].unsqueeze(-1).clone()
else:
conf = 1
if weight is not None:
weight = weight[:, None, None]
conf = conf * weight
loss = (conf * self.criterion_keypoints(pred_keypoints_2d, gt_keypoints_2d[:, :, :-1])).mean()
return loss
def keypoint_3d_loss(self, pred_keypoints_3d, gt_keypoints_3d, has_pose_3d, weight=None):
"""
Compute 3D keypoint loss for the examples that 3D keypoint annotations are available.
The loss is weighted by the weight
"""
if gt_keypoints_3d.shape[2] == 3:
tmp = gt_keypoints_3d.new_ones(gt_keypoints_3d.shape[0], gt_keypoints_3d.shape[1], 1)
gt_keypoints_3d = torch.cat((gt_keypoints_3d, tmp), dim=2)
conf = gt_keypoints_3d[:, :, -1].unsqueeze(-1).clone()
gt_keypoints_3d = gt_keypoints_3d[:, :, :-1].clone()
gt_keypoints_3d = gt_keypoints_3d[has_pose_3d == 1]
conf = conf[has_pose_3d == 1]
if weight is not None:
weight = weight[has_pose_3d == 1, None, None]
conf = conf * weight
pred_keypoints_3d = pred_keypoints_3d[has_pose_3d == 1]
if len(gt_keypoints_3d) > 0:
# Align the origin of the first 24 keypoints with the pelvis.
gt_pelvis = (gt_keypoints_3d[:, 2, :] + gt_keypoints_3d[:, 3, :]) / 2
pred_pelvis = (pred_keypoints_3d[:, 2, :] + pred_keypoints_3d[:, 3, :]) / 2
gt_keypoints_3d = gt_keypoints_3d - gt_pelvis[:, None, :]
pred_keypoints_3d = pred_keypoints_3d - pred_pelvis[:, None, :]
# # Align the origin of the first 24 keypoints with the pelvis.
# gt_pelvis = (gt_keypoints_3d[:, 2, :] + gt_keypoints_3d[:, 3, :]) / 2
# pred_pelvis = (pred_keypoints_3d[:, 2, :] + pred_keypoints_3d[:, 3, :]) / 2
# gt_keypoints_3d[:, :24, :] = gt_keypoints_3d[:, :24, :] - gt_pelvis[:, None, :]
# pred_keypoints_3d[:, :24, :] = pred_keypoints_3d[:, :24, :] - pred_pelvis[:, None, :]
#
# # Align the origin of the 24 SMPL keypoints with the root joint.
# gt_root_joint = gt_keypoints_3d[:, 24]
# pred_root_joint = pred_keypoints_3d[:, 24]
# gt_keypoints_3d[:, 24:, :] = gt_keypoints_3d[:, 24:, :] - gt_root_joint[:, None, :]
# pred_keypoints_3d[:, 24:, :] = pred_keypoints_3d[:, 24:, :] - pred_root_joint[:, None, :]
return (conf * self.criterion_keypoints_3d(pred_keypoints_3d, gt_keypoints_3d)).mean()
else:
return torch.FloatTensor(1).fill_(0.).to(self.device)
def smpl_keypoint_3d_loss(self, pred_keypoints_3d, gt_keypoints_3d, has_pose_3d, weight=None):
"""
Compute 3D SMPL keypoint loss for the examples that 3D keypoint annotations are available.
The loss is weighted by the weight
"""
if gt_keypoints_3d.shape[2] == 3:
tmp = gt_keypoints_3d.new_ones(gt_keypoints_3d.shape[0], gt_keypoints_3d.shape[1], 1)
gt_keypoints_3d = torch.cat((gt_keypoints_3d, tmp), dim=2)
conf = gt_keypoints_3d[:, :, -1].unsqueeze(-1).clone()
gt_keypoints_3d = gt_keypoints_3d[:, :, :-1].clone()
gt_keypoints_3d = gt_keypoints_3d[has_pose_3d == 1]
conf = conf[has_pose_3d == 1]
if weight is not None:
weight = weight[has_pose_3d == 1, None, None]
conf = conf * weight
pred_keypoints_3d = pred_keypoints_3d[has_pose_3d == 1]
if len(gt_keypoints_3d) > 0:
gt_root_joint = gt_keypoints_3d[:, 0, :]
pred_root_joint = pred_keypoints_3d[:, 0, :]
gt_keypoints_3d = gt_keypoints_3d - gt_root_joint[:, None, :]
pred_keypoints_3d = pred_keypoints_3d - pred_root_joint[:, None, :]
return (conf * self.criterion_keypoints_3d(pred_keypoints_3d, gt_keypoints_3d)).mean()
else:
return torch.FloatTensor(1).fill_(0.).to(self.device)
def shape_loss(self, pred_vertices, gt_vertices, has_smpl, weight=None):
"""Compute per-vertex loss on the shape for the examples that SMPL annotations are available."""
pred_vertices_with_shape = pred_vertices[has_smpl == 1]
gt_vertices_with_shape = gt_vertices[has_smpl == 1]
if weight is not None:
weight = weight[has_smpl == 1, None, None]
else:
weight = 1
if len(gt_vertices_with_shape) > 0:
loss = self.criterion_shape(pred_vertices_with_shape, gt_vertices_with_shape)
loss = (loss * weight).mean()
return loss
else:
return torch.FloatTensor(1).fill_(0.).to(self.device)
def uv_loss(self, pred_uv_map, gt_uv_map, has_smpl, weight=None):
# self.uv_mask = self.uv_mask.to(pred_uv_map.device)
self.uv_weight = self.uv_weight.to(pred_uv_map.device).type(pred_uv_map.dtype)
max = self.uv_weight.max()
pred_uv_map_shape = pred_uv_map[has_smpl == 1]
gt_uv_map_with_shape = gt_uv_map[has_smpl == 1]
if len(gt_uv_map_with_shape) > 0:
# return self.criterion_uv(pred_uv_map_shape * self.uv_mask, gt_uv_map_with_shape * self.uv_mask)
if weight is not None:
ada_weight = weight[has_smpl > 0, None, None, None]
else:
ada_weight = 1.0
loss = self.criterion_uv(pred_uv_map_shape * self.uv_weight, gt_uv_map_with_shape * self.uv_weight)
loss = (loss * ada_weight).mean()
return loss
else:
# return torch.FloatTensor(1).fill_(0.).to(self.device)
return torch.tensor(0.0, dtype=pred_uv_map.dtype, device=self.device)
def tv_loss(self, uv_map):
self.uv_weight = self.uv_weight.to(uv_map.device)
tv = torch.abs(uv_map[:,0:-1, 0:-1, :] - uv_map[:,0:-1, 1:, :]) \
+ torch.abs(uv_map[:,0:-1, 0:-1, :] - uv_map[:,1:, 0:-1, :])
return torch.sum(tv) / self.tv_factor
# return torch.sum(tv * self.uv_weight[:, 0:-1, 0:-1]) / self.tv_factor
def dp_loss(self, pred_dp, gt_dp, has_dp, weight=None):
dtype = pred_dp.dtype
pred_dp_shape = pred_dp[[has_dp > 0]]
gt_dp_shape = gt_dp[[has_dp > 0]]
if len(gt_dp_shape) > 0:
gt_mask_shape = (gt_dp_shape[:, 0].unsqueeze(1) > 0).type(dtype)
gt_uv_shape = gt_dp_shape[:, 1:]
pred_mask_shape = pred_dp_shape[:, 0].unsqueeze(1)
pred_uv_shape = pred_dp_shape[:, 1:]
pred_mask_shape = F.interpolate(pred_mask_shape, [gt_dp.shape[2], gt_dp.shape[3]], mode='bilinear')
pred_uv_shape = F.interpolate(pred_uv_shape, [gt_dp.shape[2], gt_dp.shape[3]], mode='nearest')
if weight is not None:
weight = weight[has_dp > 0, None, None, None]
else:
weight = 1.0
pred_mask_shape = pred_mask_shape.clamp(min=0.0, max=1.0)
loss_mask = torch.nn.BCELoss(reduction='none')(pred_mask_shape, gt_mask_shape)
loss_mask = (loss_mask * weight).mean()
gt_uv_weight = (gt_uv_shape.abs().max(dim=1, keepdim=True)[0] > 0).type(dtype)
weight_ratio = (gt_uv_weight.mean(dim=-1).mean(dim=-1)[:, :, None, None] + 1e-8)
gt_uv_weight = gt_uv_weight / weight_ratio # normalized the weight according to mask area
loss_uv = self.criterion_uv(gt_uv_weight * pred_uv_shape, gt_uv_weight * gt_uv_shape)
loss_uv = (loss_uv * weight).mean()
return loss_mask, loss_uv
else:
return pred_dp.sum() * 0, pred_dp.sum() * 0
def consistent_loss(self, dp, uv_map, camera, weight=None):
tmp = torch.arange(0, dp.shape[-1], 1, dtype=dp.dtype, device=dp.device) / (dp.shape[-1] -1)
tmp = tmp * 2 - 1
loc_y, loc_x = torch.meshgrid(tmp, tmp)
loc = torch.stack((loc_x, loc_y), dim=0).expand(dp.shape[0], -1, -1, -1)
dp_mask = (dp[:, 0] > 0.5).float().unsqueeze(1)
loc = dp_mask * loc
dp_tmp = dp_mask * (dp[:, 1:] * 2 - 1)
'''uv_map need to be transfered to img coordinate first'''
uv_map = uv_map[:, :, :, :-1]
camera = camera.view(-1, 1, 1, 3)
uv_map = uv_map + camera[:, :, :, 1:] # trans
uv_map = uv_map * camera[:, :, :, 0].unsqueeze(-1) # scale
warp_loc = F.grid_sample(uv_map.permute(0, 3, 1, 2), dp_tmp.permute(0, 2, 3, 1))[:, :2]
warp_loc = warp_loc * dp_mask
if weight is not None:
weight = weight[:, None, None, None]
dp_mask = dp_mask * weight
loss_con = torch.nn.MSELoss()(warp_loc * dp_mask, loc * dp_mask)
return loss_con