Skip to content

F
Fairseq-S2T
Commit f0605efa by xuchen
Options

update the pyramid transformer

parent ed623111
正在显示 包含 45 行增加 和 28 行删除
egs/librispeech/asr/conf/pyramid.yaml
...@@ -4,6 +4,7 @@ pyramid-stages: 4 ...@@ -4,6 +4,7 @@ pyramid-stages: 4
pyramid-layers: 2_2_6_2 pyramid-layers: 2_2_6_2
#encoder-attention-type: reduced #encoder-attention-type: reduced
#pyramid-attn-sample-ratios: 8_4_2_1 #pyramid-attn-sample-ratios: 8_4_2_1
#pyramid-block-attn: True
pyramid-sr-ratios: 2_2_2_2 pyramid-sr-ratios: 2_2_2_2
pyramid-use-ppm: True pyramid-use-ppm: True
pyramid-embed-dims: 128_128_256_512 pyramid-embed-dims: 128_128_256_512
...@@ -42,7 +43,7 @@ lr: 2e-3 ...@@ -42,7 +43,7 @@ lr: 2e-3
#adam_betas: (0.9,0.98) #adam_betas: (0.9,0.98)
criterion: label_smoothed_cross_entropy_with_ctc criterion: label_smoothed_cross_entropy_with_ctc
ctc-weight: 0.3 ctc-weight: 0.0
label_smoothing: 0.1 label_smoothing: 0.1
conv-channels: 1024 conv-channels: 1024
......
fairseq/criterions/label_smoothed_cross_entropy_with_ctc.py
...@@ -94,8 +94,6 @@ class LabelSmoothedCrossEntropyCriterionWithCTC( ...@@ -94,8 +94,6 @@ class LabelSmoothedCrossEntropyCriterionWithCTC(
ctc_loss = self.compute_ctc_loss(model, sample, encoder_out) ctc_loss = self.compute_ctc_loss(model, sample, encoder_out)
logging_output["ctc_loss"] = utils.item(ctc_loss.data) logging_output["ctc_loss"] = utils.item(ctc_loss.data)
loss = (1 - self.ctc_weight) * loss + self.ctc_weight * ctc_loss loss = (1 - self.ctc_weight) * loss + self.ctc_weight * ctc_loss
else:
loss = (1 - self.ctc_weight) * loss
logging_output["loss"] = utils.item(loss.data) if reduce else loss.data logging_output["loss"] = utils.item(loss.data) if reduce else loss.data
return loss, sample_size, logging_output return loss, sample_size, logging_output
......
fairseq/models/speech_to_text/pys2t_transformer.py
...@@ -55,16 +55,18 @@ class ReducedEmbed(nn.Module): ...@@ -55,16 +55,18 @@ class ReducedEmbed(nn.Module):
self.conv = nn.Conv1d(in_channels, out_channels * 2, kernel_sizes, stride=stride, padding=padding) self.conv = nn.Conv1d(in_channels, out_channels * 2, kernel_sizes, stride=stride, padding=padding)
self.glu = nn.GLU(dim=1) self.glu = nn.GLU(dim=1)
elif self.reduced_way == "proj": elif self.reduced_way == "proj":
self.proj = nn.Linear(2 * in_channels, out_channels, bias=False) self.conv = nn.Conv1d(in_channels, out_channels, kernel_sizes, padding=padding)
elif self.reduced_way == "fuse":
self.conv = nn.Conv1d(in_channels, out_channels, kernel_sizes, stride=stride, padding=padding)
self.conv_proj = nn.Conv1d(in_channels, out_channels, kernel_sizes, padding=padding)
else: else:
logger.error("Unsupported reduced way!") logger.error("Unsupported reduced way!")
self.embed_norm = embed_norm self.embed_norm = embed_norm
if self.embed_norm: if self.embed_norm:
if self.reduced_way == "proj": if self.reduced_way in ["proj", "fuse"]:
self.norm = LayerNorm(2 * in_channels) self.in_norm = nn.BatchNorm1d(in_channels)
else: self.norm = LayerNorm(out_channels)
self.norm = LayerNorm(out_channels)
def forward(self, x, lengths): def forward(self, x, lengths):
seq_len, bsz, dim = x.size() seq_len, bsz, dim = x.size()
...@@ -78,15 +80,22 @@ class ReducedEmbed(nn.Module): ...@@ -78,15 +80,22 @@ class ReducedEmbed(nn.Module):
x.masked_fill_(mask_pad, 0.0) x.masked_fill_(mask_pad, 0.0)
x = x.transpose(0, 1) x = x.transpose(0, 1)
# x = self.in_norm(x)
if self.reduced_way == "proj": if self.reduced_way == "proj":
x = x.transpose(0, 1).contiguous().view(bsz, int(seq_len / 2), -1) x = x.permute(1, 2, 0) # bsz, dim, seq_len
x = self.proj(self.norm(x)) x = nn.functional.adaptive_avg_pool1d(x, int(seq_len // self.stride))
x = x.transpose(0, 1) x = self.conv(self.in_norm(x))
x = x.permute(2, 0, 1) # seq_len, bsz, dim
else: else:
x = x.permute(1, 2, 0) # B * D * T x = x.permute(1, 2, 0) # B * D * T
origin_x = x
x = self.conv(x) x = self.conv(x)
if self.reduced_way == "glu": if self.reduced_way == "glu":
x = self.glu(x) x = self.glu(x)
if self.reduced_way == "fuse":
x2 = nn.functional.adaptive_avg_pool1d(origin_x, int(seq_len // self.stride))
x2 = self.conv_proj(self.in_norm(x2))
x = x + x2
x = x.permute(2, 0, 1) # T * B * D x = x.permute(2, 0, 1) # T * B * D
if self.embed_norm: if self.embed_norm:
x = self.norm(x) x = self.norm(x)
...@@ -182,7 +191,7 @@ class PYS2TTransformerModel(S2TTransformerModel): ...@@ -182,7 +191,7 @@ class PYS2TTransformerModel(S2TTransformerModel):
parser.add_argument( parser.add_argument(
"--pyramid-reduced-embed", "--pyramid-reduced-embed",
type=str, type=str,
choices=["glu", "conv", "proj"], choices=["glu", "conv", "proj", "fuse"],
help="the reduced way of the embedding", help="the reduced way of the embedding",
) )
parser.add_argument( parser.add_argument(
...@@ -285,12 +294,16 @@ class PyS2TTransformerEncoder(FairseqEncoder): ...@@ -285,12 +294,16 @@ class PyS2TTransformerEncoder(FairseqEncoder):
for i in range(self.pyramid_stages): for i in range(self.pyramid_stages):
num_layers = self.pyramid_layers[i] num_layers = self.pyramid_layers[i]
sr_ratio = self.pyramid_sr_ratios[i] sr_ratio = self.pyramid_sr_ratios[i]
attn_sample_ratio = self.pyramid_attn_sample_ratios[i] attn_sample_ratio = self.pyramid_attn_sample_ratios[i] if self.attn_type == "reduced" else -1
embed_dim = self.pyramid_embed_dims[i] embed_dim = self.pyramid_embed_dims[i]
kernel_size = self.pyramid_kernel_sizes[i] kernel_size = self.pyramid_kernel_sizes[i]
ffn_ratio = self.pyramid_ffn_ratios[i] ffn_ratio = self.pyramid_ffn_ratios[i]
num_head = self.pyramid_heads[i] num_head = self.pyramid_heads[i]
use_pos_embed = self.pyramid_position_embed[i] use_pos_embed = self.pyramid_position_embed[i]
logger.info("The stage {}: layer {}, sample ratio {}, attention sample ratio {}, embed dim {}, "
"kernel size {}, ffn ratio {}, num head {}, position embed {}".
format(i, num_layers, sr_ratio, attn_sample_ratio,
embed_dim, kernel_size, ffn_ratio, num_head, use_pos_embed))
if i == 0: if i == 0:
self.embed_scale = math.sqrt(embed_dim) self.embed_scale = math.sqrt(embed_dim)
...@@ -300,12 +313,11 @@ class PyS2TTransformerEncoder(FairseqEncoder): ...@@ -300,12 +313,11 @@ class PyS2TTransformerEncoder(FairseqEncoder):
reduced_embed = ReducedEmbed( reduced_embed = ReducedEmbed(
self.pyramid_reduced_embed, self.pyramid_reduced_embed,
self.pyramid_embed_norm, self.pyramid_embed_norm,
# self.pyramid_embed_norm if i != 0 else False,
args.input_feat_per_channel * args.input_channels if i == 0 else self.pyramid_embed_dims[i-1], args.input_feat_per_channel * args.input_channels if i == 0 else self.pyramid_embed_dims[i-1],
embed_dim, embed_dim,
kernel_sizes=kernel_size, kernel_sizes=kernel_size,
stride=sr_ratio, stride=sr_ratio,
padding=kernel_size // 2, padding=(kernel_size - 1) // 2,
) )
if use_pos_embed: if use_pos_embed:
pos_embed = PositionalEmbedding( pos_embed = PositionalEmbedding(
...@@ -327,13 +339,16 @@ class PyS2TTransformerEncoder(FairseqEncoder): ...@@ -327,13 +339,16 @@ class PyS2TTransformerEncoder(FairseqEncoder):
if self.use_ppm: if self.use_ppm:
ppm_layer_norm = LayerNorm(embed_dim) ppm_layer_norm = LayerNorm(embed_dim)
ppm_layer_norm2 = LayerNorm(self.embed_dim)
ppm = nn.Sequential( ppm = nn.Sequential(
nn.Conv1d(embed_dim, self.embed_dim, kernel_size=1, bias=False), nn.Conv1d(embed_dim, self.embed_dim, kernel_size=1, bias=False),
nn.BatchNorm1d(self.embed_dim), nn.BatchNorm1d(self.embed_dim),
nn.ReLU(), nn.ReLU(),
) )
else: else:
ppm_layer_norm = None ppm_layer_norm = None
ppm_layer_norm2 = None
ppm = None ppm = None
setattr(self, f"reduced_embed{i + 1}", reduced_embed) setattr(self, f"reduced_embed{i + 1}", reduced_embed)
...@@ -342,12 +357,12 @@ class PyS2TTransformerEncoder(FairseqEncoder): ...@@ -342,12 +357,12 @@ class PyS2TTransformerEncoder(FairseqEncoder):
setattr(self, f"block_attn{i + 1}", block_attn) setattr(self, f"block_attn{i + 1}", block_attn)
setattr(self, f"ppm{i + 1}", ppm) setattr(self, f"ppm{i + 1}", ppm)
setattr(self, f"ppm_layer_norm{i + 1}", ppm_layer_norm) setattr(self, f"ppm_layer_norm{i + 1}", ppm_layer_norm)
setattr(self, f"ppm_layer_norm2{i + 1}", ppm_layer_norm2)
if i == self.pyramid_stages - 1: if args.encoder_normalize_before:
if args.encoder_normalize_before: self.layer_norm = LayerNorm(self.embed_dim)
self.layer_norm = LayerNorm(embed_dim) else:
else: self.layer_norm = None
self.layer_norm = None
if self.use_ppm: if self.use_ppm:
self.ppm_weight = nn.Parameter(torch.Tensor(self.pyramid_stages).fill_(1.0)) self.ppm_weight = nn.Parameter(torch.Tensor(self.pyramid_stages).fill_(1.0))
...@@ -442,21 +457,25 @@ class PyS2TTransformerEncoder(FairseqEncoder): ...@@ -442,21 +457,25 @@ class PyS2TTransformerEncoder(FairseqEncoder):
x = block_attn(x, prev_state[-1], prev_padding[-1]) x = block_attn(x, prev_state[-1], prev_padding[-1])
if self.use_ppm: if self.use_ppm:
pool_state = [x] pool_state = []
seq_len, bsz, dim = x.size() seq_len, bsz, dim = x.size()
i = -1 i = -1
for state in prev_state[:-1]: for state in prev_state:
i += 1 i += 1
ppm = getattr(self, f"ppm{i + 1}") ppm = getattr(self, f"ppm{i + 1}")
ppm_layer_norm = getattr(self, f"ppm_layer_norm{i + 1}") ppm_layer_norm = getattr(self, f"ppm_layer_norm{i + 1}")
ppm_layer_norm2 = getattr(self, f"ppm_layer_norm2{i + 1}")
state = ppm_layer_norm(state) state = ppm_layer_norm(state)
state = state.permute(1, 2, 0) state = state.permute(1, 2, 0) # bsz, dim, seq_len
state = nn.functional.adaptive_avg_pool1d(state, seq_len) if i != self.pyramid_stages - 1:
state = nn.functional.adaptive_avg_pool1d(state, seq_len)
state = ppm(state) state = ppm(state)
state = state.permute(2, 0, 1) state = state.permute(2, 0, 1)
state = ppm_layer_norm2(state)
pool_state.append(state) pool_state.append(state)
x = (torch.stack(pool_state, dim=0) * self.ppm_weight.view(-1, 1, 1, 1)).sum(0) ppm_weight = self.ppm_weight
x = (torch.stack(pool_state, dim=0) * ppm_weight.view(-1, 1, 1, 1)).sum(0)
if self.layer_norm is not None: if self.layer_norm is not None:
x = self.layer_norm(x) x = self.layer_norm(x)
......
fairseq/modules/rel_position_multihead_attention.py
...@@ -128,7 +128,6 @@ class RelPositionMultiheadAttention(MultiheadAttention): ...@@ -128,7 +128,6 @@ class RelPositionMultiheadAttention(MultiheadAttention):
else: else:
k = self.k_proj(key) k = self.k_proj(key)
v = self.v_proj(key) v = self.v_proj(key)
else: else:
assert key is not None and value is not None assert key is not None and value is not None
q = self.q_proj(query) q = self.q_proj(query)
...@@ -301,7 +300,7 @@ class RelPositionMultiheadAttention(MultiheadAttention): ...@@ -301,7 +300,7 @@ class RelPositionMultiheadAttention(MultiheadAttention):
x = x * torch.tril(ones, x.size(2) - x.size(1))[None, :, :] x = x * torch.tril(ones, x.size(2) - x.size(1))[None, :, :]
return x return x
# matrix_bd = rel_shift(matrix_bd) matrix_bd = rel_shift(matrix_bd)
attn_weights = (matrix_ac + matrix_bd) * self.scaling attn_weights = (matrix_ac + matrix_bd) * self.scaling
...@@ -456,7 +455,7 @@ class RelPositionMultiheadAttention(MultiheadAttention): ...@@ -456,7 +455,7 @@ class RelPositionMultiheadAttention(MultiheadAttention):
dim = int(state_dict[k].shape[0] / 3) dim = int(state_dict[k].shape[0] / 3)
items_to_add[prefix + "q_proj.bias"] = state_dict[k_bias][:dim] items_to_add[prefix + "q_proj.bias"] = state_dict[k_bias][:dim]
items_to_add[prefix + "k_proj.bias"] = state_dict[k_bias][ items_to_add[prefix + "k_proj.bias"] = state_dict[k_bias][
dim : 2 * dim dim: 2 * dim
] ]
items_to_add[prefix + "v_proj.bias"] = state_dict[k_bias][2 * dim :] items_to_add[prefix + "v_proj.bias"] = state_dict[k_bias][2 * dim :]
......
Markdown 格式
0% 或
您添加了 0 人 到此讨论。请谨慎行事。
请先完成此评论的编辑！
请 注册 或者 后发表评论