def get_criterion(conf, device="cuda"):
criterion = {
"mse": nn.MSELoss(),
"l1": nn.L1Loss(),
"ce": nn.CrossEntropyLoss(ignore_index=-100),
"kld": nn.KLDivLoss(reduction="mean"),
"fmse": CustomFeatureLoss(loss_type="mse", causal_size=conf["causal_size"]),
"fl1": CustomFeatureLoss(loss_type="l1", causal_size=conf["causal_size"]),
"fstft": CustomFeatureLoss(
loss_type="stft",
causal_size=conf["causal_size"],
stft_params=conf["stft_params"],
device=device,
),
}
return criterion
def test_customloss():
x = torch.randn((B, T, D))
y = torch.randn((B, T, D))
mask = x.ge(0)
for c in [-8, -2, 0, 2, 8]:
for loss_type in ["l1", "mse", "stft"]:
criterion = CustomFeatureLoss(loss_type=loss_type, causal_size=c)
if loss_type != "stft":
_ = criterion(x, y, mask=mask)
else:
_ = criterion(x, y)