File size: 3,513 Bytes
4409449 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 |
import torch
import torch.nn as nn
from .base import BaseLosses
class CommitLoss(nn.Module):
"""
Useless Wrapper
"""
def __init__(self, **kwargs):
super().__init__()
def forward(self, commit, commit2, **kwargs):
return commit
class GPTLosses(BaseLosses):
def __init__(self, cfg, stage, num_joints, **kwargs):
# Save parameters
self.stage = stage
recons_loss = cfg.LOSS.ABLATION.RECONS_LOSS
# Define losses
losses = []
params = {}
if stage == "vae":
losses.append("recons_feature")
params['recons_feature'] = cfg.LOSS.LAMBDA_FEATURE
losses.append("recons_velocity")
params['recons_velocity'] = cfg.LOSS.LAMBDA_VELOCITY
losses.append("vq_commit")
params['vq_commit'] = cfg.LOSS.LAMBDA_COMMIT
elif stage in ["lm_pretrain", "lm_instruct"]:
losses.append("gpt_loss")
params['gpt_loss'] = cfg.LOSS.LAMBDA_CLS
# Define loss functions & weights
losses_func = {}
for loss in losses:
if loss.split('_')[0] == 'recons':
if recons_loss == "l1":
losses_func[loss] = nn.L1Loss
elif recons_loss == "l2":
losses_func[loss] = nn.MSELoss
elif recons_loss == "l1_smooth":
losses_func[loss] = nn.SmoothL1Loss
elif loss.split('_')[1] in [
'commit', 'loss', 'gpt', 'm2t2m', 't2m2t'
]:
losses_func[loss] = CommitLoss
elif loss.split('_')[1] in ['cls', 'lm']:
losses_func[loss] = nn.CrossEntropyLoss
else:
raise NotImplementedError(f"Loss {loss} not implemented.")
super().__init__(cfg, losses, params, losses_func, num_joints,
**kwargs)
def update(self, rs_set):
'''Update the losses'''
total: float = 0.0
if self.stage in ["vae"]:
total += self._update_loss("recons_feature", rs_set['m_rst'],
rs_set['m_ref'])
# total += self._update_loss("recons_joints", rs_set['joints_rst'], rs_set['joints_ref'])
nfeats = rs_set['m_rst'].shape[-1]
if nfeats in [263, 135 + 263]:
if nfeats == 135 + 263:
vel_start = 135 + 4
elif nfeats == 263:
vel_start = 4
total += self._update_loss(
"recons_velocity",
rs_set['m_rst'][..., vel_start:(self.num_joints - 1) * 3 +
vel_start],
rs_set['m_ref'][..., vel_start:(self.num_joints - 1) * 3 +
vel_start])
else:
if self._params['recons_velocity'] != 0.0:
raise NotImplementedError(
"Velocity not implemented for nfeats = {})".format(nfeats))
total += self._update_loss("vq_commit", rs_set['loss_commit'],
rs_set['loss_commit'])
if self.stage in ["lm_pretrain", "lm_instruct"]:
total += self._update_loss("gpt_loss", rs_set['outputs'].loss,
rs_set['outputs'].loss)
# Update the total loss
self.total += total.detach()
self.count += 1
return total
|