diff --git a/layers/losses.py b/layers/losses.py index f2b6fcb0..a9d393cb 100644 --- a/layers/losses.py +++ b/layers/losses.py @@ -17,10 +17,10 @@ def _sequence_mask(sequence_length, max_len=None): return seq_range_expand < seq_length_expand -class L1LossMasked(nn.Module): +class L2LossMasked(nn.Module): def __init__(self): - super(L1LossMasked, self).__init__() + super(L2LossMasked, self).__init__() def forward(self, input, target, length): """ @@ -44,7 +44,7 @@ class L1LossMasked(nn.Module): # target_flat: (batch * max_len, dim) target_flat = target.view(-1, target.shape[-1]) # losses_flat: (batch * max_len, dim) - losses_flat = functional.l1_loss(input, target_flat, size_average=False, + losses_flat = functional.mse_loss(input, target_flat, size_average=False, reduce=False) # losses: (batch, max_len, dim) losses = losses_flat.view(*target.size()) diff --git a/train.py b/train.py index f570b7fe..954c300e 100644 --- a/train.py +++ b/train.py @@ -26,7 +26,7 @@ from utils.model import get_param_size from utils.visual import plot_alignment, plot_spectrogram from datasets.LJSpeech import LJSpeechDataset from models.tacotron import Tacotron -from layers.losses import L1LossMasked +from layers.losses import L2LossMasked torch.manual_seed(1) use_cuda = torch.cuda.is_available() @@ -365,7 +365,7 @@ def main(args): optimizer = optim.Adam(model.parameters(), lr=c.lr) optimizer_st = optim.Adam(model.decoder.stopnet.parameters(), lr=c.lr) - criterion = L1LossMasked() + criterion = L2LossMasked() criterion_st = nn.BCELoss() if args.restore_path: