chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,9 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
from . import multilingual_translation_latent_depth # noqa
|
||||
from .loss import latent_depth # noqa
|
||||
from .models import latent_multilingual_transformer # noqa
|
||||
from .modules import latent_layers # noqa
|
||||
@@ -0,0 +1,99 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import math
|
||||
|
||||
import torch
|
||||
from torch.nn.modules.loss import _Loss
|
||||
|
||||
|
||||
class LatentLayersKLLoss(_Loss):
|
||||
def __init__(self, args):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
|
||||
def forward(self, layer_samples, lang_idx, update_num, sample_size):
|
||||
prior = self.args.prior
|
||||
samples = layer_samples[lang_idx]
|
||||
eps = 1e-7
|
||||
if prior == "uniform":
|
||||
# uniform prior
|
||||
kl_loss = (samples * (torch.log(samples + eps) - math.log(0.5))).sum(-1)
|
||||
elif prior == "agged_posterior":
|
||||
# aggregated posterior
|
||||
y_t = torch.stack([x.detach() for x in layer_samples], dim=0)
|
||||
agged_q = torch.sum(y_t, dim=0)
|
||||
row_norm = agged_q.sum(-1)
|
||||
normed_agg_q = agged_q / row_norm
|
||||
kl_loss = (
|
||||
samples * (torch.log(samples + eps) - torch.log(normed_agg_q + eps))
|
||||
).sum(-1)
|
||||
else:
|
||||
raise NotImplementedError("The specified prior is not implemented.")
|
||||
|
||||
# normalized by number of layers
|
||||
kl_loss /= layer_samples[0].size()[0]
|
||||
kl_weight = min(
|
||||
self.args.sparsity_weight,
|
||||
(update_num - self.args.soft_update)
|
||||
* self.args.sparsity_weight
|
||||
/ self.args.anneal_updates,
|
||||
)
|
||||
kl_loss *= kl_weight * sample_size
|
||||
return kl_loss
|
||||
|
||||
|
||||
class LatentLayersSparsityLoss(_Loss):
|
||||
def __init__(self, args):
|
||||
super().__init__()
|
||||
self.args = args
|
||||
|
||||
def is_valid(self, update_num):
|
||||
if self.args.target_layers <= 0:
|
||||
return False
|
||||
return update_num > (self.args.soft_update + self.args.anneal_updates)
|
||||
|
||||
def forward(self, layer_samples_list, update_num, sample_size):
|
||||
batch_loss = 0
|
||||
share_loss = 0
|
||||
global_sparsity_loss = 0
|
||||
layer_samples = torch.stack(layer_samples_list, dim=0)
|
||||
if (
|
||||
self.args.target_layers > 0 or self.args.share_weight > 0
|
||||
) and update_num > (self.args.soft_update + self.args.anneal_updates):
|
||||
# anneal sparsity weight
|
||||
if update_num < (self.args.anneal_updates + self.args.soft_update):
|
||||
weight_anneal = 0
|
||||
elif update_num < (2 * self.args.anneal_updates + self.args.soft_update):
|
||||
weight_anneal = (
|
||||
(update_num - self.args.soft_update - self.args.anneal_updates)
|
||||
* self.args.share_weight
|
||||
/ self.args.anneal_updates
|
||||
)
|
||||
else:
|
||||
weight_anneal = 1
|
||||
# compute ratio among languages
|
||||
layer_utilization = torch.sum(layer_samples, dim=0)
|
||||
layer_utilization /= layer_samples.size()[0]
|
||||
if self.args.share_weight > 0:
|
||||
# encouraging sharing across languages
|
||||
share_loss = sum(
|
||||
-1.0 * v * math.log(v) for v in layer_utilization if v > 0
|
||||
)
|
||||
batch_loss += (
|
||||
weight_anneal * self.args.share_weight * sample_size * share_loss
|
||||
)
|
||||
if self.args.target_layers > 0:
|
||||
# computed expected number of layers selected
|
||||
expeted_layers = sum(layer_utilization)
|
||||
# compute l2 loss wrt target number of layers
|
||||
global_sparsity_loss = (expeted_layers - self.args.target_layers) ** 2
|
||||
batch_loss += (
|
||||
weight_anneal
|
||||
* self.args.share_weight
|
||||
* sample_size
|
||||
* global_sparsity_loss
|
||||
)
|
||||
return batch_loss
|
||||
+75
@@ -0,0 +1,75 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
from fairseq.models import register_model, register_model_architecture
|
||||
from fairseq.models.multilingual_transformer import MultilingualTransformerModel
|
||||
from fairseq.models.transformer import (
|
||||
TransformerDecoder,
|
||||
TransformerEncoder,
|
||||
base_architecture,
|
||||
)
|
||||
|
||||
from .latent_transformer import LatentTransformerDecoder, LatentTransformerEncoder
|
||||
|
||||
|
||||
@register_model("latent_multilingual_transformer")
|
||||
class LatentMultilingualTransformerModel(MultilingualTransformerModel):
|
||||
"""A variant of standard multilingual Transformer models which encoder and/or
|
||||
decoders supports latent depth, as is in "Deep Transformer with Latent Depth"
|
||||
(https://arxiv.org/abs/2009.13102).
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def add_args(parser):
|
||||
"""Add model-specific arguments to the parser."""
|
||||
MultilingualTransformerModel.add_args(parser)
|
||||
parser.add_argument(
|
||||
'--soft-select',
|
||||
action='store_true',
|
||||
help='use soft samples in training an inference',
|
||||
)
|
||||
parser.add_argument(
|
||||
'--sampling-tau',
|
||||
type=float,
|
||||
default=5.,
|
||||
help='sampling temperature',
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _get_module_class(cls, is_encoder, args, lang_dict, embed_tokens, langs):
|
||||
if is_encoder:
|
||||
if hasattr(args, "encoder_latent_layer") and args.encoder_latent_layer:
|
||||
return LatentTransformerEncoder(
|
||||
args, lang_dict, embed_tokens, num_logits=len(langs)
|
||||
)
|
||||
else:
|
||||
return TransformerEncoder(args, lang_dict, embed_tokens)
|
||||
else:
|
||||
if hasattr(args, "decoder_latent_layer") and args.decoder_latent_layer:
|
||||
return LatentTransformerDecoder(
|
||||
args, lang_dict, embed_tokens, num_logits=len(langs)
|
||||
)
|
||||
else:
|
||||
return TransformerDecoder(args, lang_dict, embed_tokens)
|
||||
|
||||
|
||||
@register_model_architecture(
|
||||
"latent_multilingual_transformer", "latent_multilingual_transformer"
|
||||
)
|
||||
def latent_multilingual_architecture(args):
|
||||
args.encoder_embed_dim = getattr(args, "encoder_embed_dim", 512)
|
||||
args.encoder_ffn_embed_dim = getattr(args, "encoder_ffn_embed_dim", 1024)
|
||||
args.encoder_attention_heads = getattr(args, "encoder_attention_heads", 4)
|
||||
args.encoder_layers = getattr(args, "encoder_layers", 12)
|
||||
args.decoder_embed_dim = getattr(args, "decoder_embed_dim", 512)
|
||||
args.decoder_ffn_embed_dim = getattr(args, "decoder_ffn_embed_dim", 1024)
|
||||
args.decoder_attention_heads = getattr(args, "decoder_attention_heads", 4)
|
||||
args.decoder_layers = getattr(args, "decoder_layers", 24)
|
||||
args.share_encoders = getattr(args, "share_encoders", True)
|
||||
args.share_decoders = getattr(args, "share_decoders", True)
|
||||
args.share_encoder_embeddings = getattr(args, "share_encoder_embeddings", True)
|
||||
args.share_decoder_embeddings = getattr(args, "share_decoder_embeddings", True)
|
||||
|
||||
base_architecture(args)
|
||||
+156
@@ -0,0 +1,156 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
import torch.nn as nn
|
||||
from fairseq.models.fairseq_encoder import EncoderOut
|
||||
from fairseq.models.transformer import TransformerDecoder, TransformerEncoder
|
||||
from fairseq.modules import TransformerDecoderLayer, TransformerEncoderLayer
|
||||
from torch import Tensor
|
||||
|
||||
from ..modules.latent_layers import LayerSelect
|
||||
|
||||
|
||||
class LatentTransformerEncoder(TransformerEncoder):
|
||||
"""Latent depth (https://arxiv.org/abs/2009.13102) implemented in
|
||||
TransformerEncoder.
|
||||
"""
|
||||
|
||||
def __init__(self, args, dictionary, embed_tokens, num_logits=1):
|
||||
self.num_logits = num_logits
|
||||
self.num_layers = args.encoder_layers
|
||||
super().__init__(args, dictionary, embed_tokens)
|
||||
self.layer_select = LayerSelect(
|
||||
num_layers=self.num_layers,
|
||||
num_logits=self.num_logits,
|
||||
soft_select=getattr(args, "soft_select", False),
|
||||
sampling_tau=getattr(args, "sampling_tau", 5.),
|
||||
)
|
||||
self.lang_idx = None
|
||||
self.layers = nn.ModuleList(
|
||||
[self._build_encoder_layer(args, idx) for idx in range(args.encoder_layers)]
|
||||
)
|
||||
|
||||
def set_lang_idx(self, lang_idx):
|
||||
self.lang_idx = lang_idx
|
||||
|
||||
def _build_encoder_layer(self, args, idx=None):
|
||||
return LatentTransformerEncoderLayer(args, idx, layer_select=self.layer_select)
|
||||
|
||||
def forward(self, src_tokens, src_lengths, return_all_hiddens: bool = False):
|
||||
self.layer_select.sample(self.lang_idx)
|
||||
return super().forward(src_tokens, src_lengths, return_all_hiddens)
|
||||
|
||||
|
||||
class LatentTransformerEncoderLayer(TransformerEncoderLayer):
|
||||
"""Encoder layer with each (non_residual) block weighted by samples of Bernouli
|
||||
or Gumbel Signmoid samples.
|
||||
|
||||
Args:
|
||||
args (argparse.Namespace): parsed command-line arguments from standard
|
||||
TransformerEncoderLayer.
|
||||
idx (int): layer index (used to retrieve samples).
|
||||
layer_select (LayerSelect, optional): instance of LayerSelect module with logits
|
||||
parameters and sampling method.
|
||||
"""
|
||||
|
||||
def __init__(self, args, idx, layer_select=None):
|
||||
super().__init__(args)
|
||||
self.idx = idx
|
||||
self.layer_select = layer_select
|
||||
|
||||
def residual_connection(self, x, residual):
|
||||
return residual + x * self.layer_select(self.idx)
|
||||
|
||||
|
||||
class LatentTransformerDecoder(TransformerDecoder):
|
||||
"""Latent depth (https://arxiv.org/abs/2009.13102) implemented in
|
||||
TransformerDecoder.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, args, dictionary, embed_tokens, no_encoder_attn=False, num_logits=1
|
||||
):
|
||||
self.num_logits = num_logits
|
||||
self.num_layers = args.decoder_layers
|
||||
super().__init__(
|
||||
args, dictionary, embed_tokens, no_encoder_attn=no_encoder_attn
|
||||
)
|
||||
self.layer_select = LayerSelect(
|
||||
num_layers=self.num_layers,
|
||||
num_logits=self.num_logits,
|
||||
soft_select=getattr(args, "soft_select", False),
|
||||
sampling_tau=getattr(args, "sampling_tau", 5.),
|
||||
)
|
||||
self.lang_idx = None
|
||||
self.layers = nn.ModuleList(
|
||||
[
|
||||
self._build_decoder_layer(args, no_encoder_attn, idx)
|
||||
for idx in range(args.decoder_layers)
|
||||
]
|
||||
)
|
||||
|
||||
def set_lang_idx(self, lang_idx):
|
||||
self.lang_idx = lang_idx
|
||||
|
||||
def _build_decoder_layer(self, args, no_encoder_attn=False, idx=None):
|
||||
return LatentTransformerDecoderLayer(
|
||||
args, idx, layer_select=self.layer_select, no_encoder_attn=no_encoder_attn
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
prev_output_tokens,
|
||||
encoder_out: Optional[EncoderOut] = None,
|
||||
incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]] = None,
|
||||
features_only: bool = False,
|
||||
alignment_layer: Optional[int] = None,
|
||||
alignment_heads: Optional[int] = None,
|
||||
src_lengths: Optional[Any] = None,
|
||||
return_all_hiddens: bool = False,
|
||||
):
|
||||
self.layer_select.sample(self.lang_idx)
|
||||
return super().forward(
|
||||
prev_output_tokens=prev_output_tokens,
|
||||
encoder_out=encoder_out,
|
||||
incremental_state=incremental_state,
|
||||
features_only=features_only,
|
||||
alignment_layer=alignment_layer,
|
||||
src_lengths=src_lengths,
|
||||
return_all_hiddens=return_all_hiddens,
|
||||
)
|
||||
|
||||
|
||||
class LatentTransformerDecoderLayer(TransformerDecoderLayer):
|
||||
"""Decoder layer with each (non_residual) block weighted by samples of Bernouli
|
||||
or Gumbel Signmoid samples.
|
||||
|
||||
Args:
|
||||
args (argparse.Namespace): parsed command-line arguments from standard
|
||||
TransformerDecoderLayer.
|
||||
idx (int): layer index (used to retrieve samples).
|
||||
layer_select (LayerSelect, optional): instance of LayerSelect module with logits
|
||||
parameters and sampling method.
|
||||
no_encoder_attn (bool, optional): whether to attend to encoder outputs
|
||||
(default: False).
|
||||
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
args,
|
||||
idx,
|
||||
layer_select=None,
|
||||
no_encoder_attn=False,
|
||||
add_bias_kv=False,
|
||||
add_zero_attn=False,
|
||||
):
|
||||
super().__init__(args, no_encoder_attn, add_bias_kv, add_zero_attn)
|
||||
self.idx = idx
|
||||
self.layer_select = layer_select
|
||||
|
||||
def residual_connection(self, x, residual):
|
||||
return residual + x * self.layer_select(self.idx)
|
||||
@@ -0,0 +1,75 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class LayerSelect(nn.Module):
|
||||
"""Compute samples (from a Gumbel-Sigmoid distribution) which is used as
|
||||
either (soft) weighting or (hard) selection of residual connection.
|
||||
https://arxiv.org/abs/2009.13102
|
||||
"""
|
||||
def __init__(self, num_layers, num_logits, soft_select=False, sampling_tau=5.):
|
||||
super(LayerSelect, self).__init__()
|
||||
self.layer_logits = torch.nn.Parameter(
|
||||
torch.Tensor(num_logits, num_layers),
|
||||
requires_grad=True,
|
||||
)
|
||||
self.hard_select = not soft_select
|
||||
self.tau = sampling_tau
|
||||
self.detach_grad = False
|
||||
self.layer_samples = [None] * num_logits
|
||||
|
||||
def sample(self, logit_idx):
|
||||
"""To leverage the efficiency of distributed training, samples for all
|
||||
layers are computed at once for each logit_idx. Logits are parameters
|
||||
learnt independent of each other.
|
||||
|
||||
Args:
|
||||
logit_idx: The index of logit parameters used for sampling.
|
||||
"""
|
||||
assert logit_idx is not None
|
||||
self.samples = self._gumbel_sigmoid(
|
||||
self.layer_logits[logit_idx, :].detach()
|
||||
if self.detach_grad
|
||||
else self.layer_logits[logit_idx, :],
|
||||
dim=-1,
|
||||
tau=self.tau,
|
||||
hard=self.hard_select,
|
||||
)
|
||||
self.layer_samples[logit_idx] = self.samples
|
||||
|
||||
def forward(self, i):
|
||||
sample = self.samples[i]
|
||||
return sample
|
||||
|
||||
def _gumbel_sigmoid(
|
||||
self, logits, tau=1, hard=False, eps=1e-10, dim=-1, threshold=0.5
|
||||
):
|
||||
# ~Gumbel(0,1)
|
||||
gumbels1 = (
|
||||
-torch.empty_like(logits, memory_format=torch.legacy_contiguous_format)
|
||||
.exponential_()
|
||||
.log()
|
||||
)
|
||||
gumbels2 = (
|
||||
-torch.empty_like(logits, memory_format=torch.legacy_contiguous_format)
|
||||
.exponential_()
|
||||
.log()
|
||||
)
|
||||
# Difference of two gumbels because we apply a sigmoid
|
||||
gumbels1 = (logits + gumbels1 - gumbels2) / tau
|
||||
y_soft = gumbels1.sigmoid()
|
||||
if hard:
|
||||
# Straight through.
|
||||
y_hard = torch.zeros_like(
|
||||
logits, memory_format=torch.legacy_contiguous_format
|
||||
).masked_fill(y_soft > threshold, 1.0)
|
||||
ret = y_hard - y_soft.detach() + y_soft
|
||||
else:
|
||||
# Reparametrization trick.
|
||||
ret = y_soft
|
||||
return ret
|
||||
+194
@@ -0,0 +1,194 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
from fairseq.tasks import register_task
|
||||
from fairseq.tasks.multilingual_translation import MultilingualTranslationTask
|
||||
|
||||
from .loss.latent_depth import LatentLayersKLLoss, LatentLayersSparsityLoss
|
||||
|
||||
|
||||
@register_task("multilingual_translation_latent_depth")
|
||||
class MultilingualTranslationTaskLatentDepth(MultilingualTranslationTask):
|
||||
"""A task for multiple translation with latent depth.
|
||||
|
||||
See `"Deep Transformer with Latent Depth"
|
||||
(Li et al., 2020) <https://arxiv.org/pdf/2009.13102.pdf>`_.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def add_args(parser):
|
||||
"""Add task-specific arguments to the parser."""
|
||||
# fmt: off
|
||||
MultilingualTranslationTask.add_args(parser)
|
||||
parser.add_argument('--encoder-latent-layer', action='store_true', help='latent layer selection in encoder')
|
||||
parser.add_argument('--decoder-latent-layer', action='store_true', help='latent layer selection in decoder')
|
||||
parser.add_argument('--target-layers', default=-1, type=int,
|
||||
help='number of effective layers to learn; -1 means no constraint')
|
||||
parser.add_argument('--sparsity-weight', default=0.0, type=float,
|
||||
help='weight for sparsity loss')
|
||||
parser.add_argument('--share-weight', default=0.0, type=float,
|
||||
help='weight for sharing loss')
|
||||
parser.add_argument('--soft-update', default=1, type=int,
|
||||
help='number of updates with soft sampling')
|
||||
parser.add_argument('--anneal-updates', default=1, type=int,
|
||||
help='number of updates to anneal the KL loss weight')
|
||||
parser.add_argument('--prior', default="uniform", type=str,
|
||||
help='prior used for computing KL loss')
|
||||
# fmt: on
|
||||
|
||||
def __init__(self, args, dicts, training):
|
||||
super().__init__(args, dicts, training)
|
||||
self.src_langs, self.tgt_langs = zip(
|
||||
*[(lang.split("-")[0], lang.split("-")[1]) for lang in args.lang_pairs]
|
||||
)
|
||||
if self.training and self.encoder_latent_layer:
|
||||
assert self.args.share_encoders
|
||||
if self.training and self.decoder_latent_layer:
|
||||
assert self.args.share_decoders
|
||||
if training or self.encoder_latent_layer or self.decoder_latent_layer:
|
||||
self.lang_pairs = args.lang_pairs
|
||||
else:
|
||||
self.lang_pairs = ["{}-{}".format(args.source_lang, args.target_lang)]
|
||||
self.eval_lang_pairs = self.lang_pairs
|
||||
self.model_lang_pairs = self.lang_pairs
|
||||
if self.training and (self.encoder_latent_layer or self.decoder_latent_layer):
|
||||
self.kl_loss = LatentLayersKLLoss(self.args)
|
||||
self.sparsity_loss = LatentLayersSparsityLoss(self.args)
|
||||
|
||||
def _per_lang_pair_train_loss(
|
||||
self, lang_pair, model, update_num, criterion, sample, optimizer, ignore_grad
|
||||
):
|
||||
src, tgt = lang_pair.split("-")
|
||||
if self.encoder_latent_layer:
|
||||
src_lang_idx = self.src_lang_idx_dict[src]
|
||||
model.models[lang_pair].encoder.set_lang_idx(src_lang_idx)
|
||||
model.models[lang_pair].encoder.layer_select.hard_select = (
|
||||
update_num > self.args.soft_update
|
||||
)
|
||||
if self.decoder_latent_layer:
|
||||
tgt_lang_idx = self.tgt_lang_idx_dict[tgt]
|
||||
model.models[lang_pair].decoder.set_lang_idx(tgt_lang_idx)
|
||||
model.models[lang_pair].decoder.layer_select.hard_select = (
|
||||
update_num > self.args.soft_update
|
||||
)
|
||||
|
||||
loss, sample_size, logging_output = criterion(
|
||||
model.models[lang_pair], sample[lang_pair]
|
||||
)
|
||||
if self.encoder_latent_layer:
|
||||
none_samples = sum(
|
||||
1 if x is None else 0
|
||||
for x in model.models[lang_pair].encoder.layer_select.layer_samples
|
||||
)
|
||||
if none_samples == 0 or self.args.prior != "agged_posterior":
|
||||
loss += self.kl_loss(
|
||||
model.models[lang_pair].encoder.layer_select.layer_samples,
|
||||
src_lang_idx,
|
||||
update_num,
|
||||
sample_size,
|
||||
)
|
||||
if self.decoder_latent_layer:
|
||||
none_samples = sum(
|
||||
1 if x is None else 0
|
||||
for x in model.models[lang_pair].decoder.layer_select.layer_samples
|
||||
)
|
||||
if none_samples == 0 or self.args.prior != "agged_posterior":
|
||||
loss += self.kl_loss(
|
||||
model.models[lang_pair].decoder.layer_select.layer_samples,
|
||||
tgt_lang_idx,
|
||||
update_num,
|
||||
sample_size,
|
||||
)
|
||||
if ignore_grad:
|
||||
loss *= 0
|
||||
|
||||
if hasattr(self, "sparsity_loss") and self.sparsity_loss.is_valid(update_num):
|
||||
# need to retain the graph if sparsity loss needs to be added
|
||||
loss.backward(retain_graph=True)
|
||||
else:
|
||||
optimizer.backward(loss)
|
||||
|
||||
return loss, sample_size, logging_output
|
||||
|
||||
def train_step(
|
||||
self, sample, model, criterion, optimizer, update_num, ignore_grad=False
|
||||
):
|
||||
agg_loss, agg_sample_size, agg_logging_output = super().train_step(
|
||||
sample, model, criterion, optimizer, update_num, ignore_grad
|
||||
)
|
||||
# compute auxiliary loss from layere sparsity, based on all samples from all languages
|
||||
if hasattr(self, "sparsity_loss") and self.sparsity_loss.is_valid(update_num):
|
||||
sparsity_loss = 0
|
||||
if self.encoder_latent_layer:
|
||||
sparsity_loss += self.sparsity_loss(
|
||||
next(
|
||||
iter(model.models.values())
|
||||
).encoder.layer_select.layer_samples,
|
||||
update_num,
|
||||
agg_sample_size,
|
||||
)
|
||||
if self.decoder_latent_layer:
|
||||
sparsity_loss += self.sparsity_loss(
|
||||
next(
|
||||
iter(model.models.values())
|
||||
).decoder.layer_select.layer_samples,
|
||||
update_num,
|
||||
agg_sample_size,
|
||||
)
|
||||
if sparsity_loss > 0:
|
||||
optimizer.backward(sparsity_loss)
|
||||
return agg_loss, agg_sample_size, agg_logging_output
|
||||
|
||||
def _per_lang_pair_valid_loss(self, lang_pair, model, criterion, sample):
|
||||
src, tgt = lang_pair.split("-")
|
||||
if self.encoder_latent_layer:
|
||||
src_lang_idx = self.src_lang_idx_dict[src]
|
||||
model.models[lang_pair].encoder.set_lang_idx(src_lang_idx)
|
||||
if self.decoder_latent_layer:
|
||||
tgt_lang_idx = self.tgt_lang_idx_dict[tgt]
|
||||
model.models[lang_pair].decoder.set_lang_idx(tgt_lang_idx)
|
||||
loss, sample_size, logging_output = criterion(
|
||||
model.models[lang_pair], sample[lang_pair]
|
||||
)
|
||||
return loss, sample_size, logging_output
|
||||
|
||||
def inference_step(
|
||||
self, generator, models, sample, prefix_tokens=None, constraints=None
|
||||
):
|
||||
if self.encoder_latent_layer or self.decoder_latent_layer:
|
||||
for model in models:
|
||||
if self.encoder_latent_layer:
|
||||
assert model.encoder.layer_select is not None
|
||||
src_lang_idx = self.src_lang_idx_dict[self.args.source_lang]
|
||||
model.encoder.set_lang_idx(src_lang_idx)
|
||||
if self.decoder_latent_layer:
|
||||
assert model.decoder.layer_select is not None
|
||||
tgt_lang_idx = self.tgt_lang_idx_dict[self.args.target_lang]
|
||||
model.decoder.set_lang_idx(tgt_lang_idx)
|
||||
return super().inference_step(
|
||||
generator, models, sample, prefix_tokens, constraints
|
||||
)
|
||||
|
||||
@property
|
||||
def encoder_latent_layer(self):
|
||||
return (
|
||||
hasattr(self.args, "encoder_latent_layer")
|
||||
and self.args.encoder_latent_layer
|
||||
)
|
||||
|
||||
@property
|
||||
def decoder_latent_layer(self):
|
||||
return (
|
||||
hasattr(self.args, "decoder_latent_layer")
|
||||
and self.args.decoder_latent_layer
|
||||
)
|
||||
|
||||
@property
|
||||
def src_lang_idx_dict(self):
|
||||
return {lang: lang_idx for lang_idx, lang in enumerate(self.src_langs)}
|
||||
|
||||
@property
|
||||
def tgt_lang_idx_dict(self):
|
||||
return {lang: lang_idx for lang_idx, lang in enumerate(self.tgt_langs)}
|
||||
Reference in New Issue
Block a user