Add AudioCaps checkpoint

app.py

@@ -0,0 +1,97 @@
+from pathlib import Path
+import argparse
+from functools import partial
+import gradio as gr
+import torch
+from torchaudio.functional import resample
+
+import utils.train_util as train_util
+
+
+def load_model(cfg,
+               ckpt_path,
+               device):
+    model = train_util.init_model_from_config(cfg["model"])
+    ckpt = torch.load(ckpt_path, "cpu")
+    train_util.load_pretrained_model(model, ckpt)
+    model.eval()
+    model = model.to(device)
+    tokenizer = train_util.init_obj_from_dict(cfg["tokenizer"])
+    if not tokenizer.loaded:
+        tokenizer.load_state_dict(ckpt["tokenizer"])
+    model.set_index(tokenizer.bos, tokenizer.eos, tokenizer.pad)
+    return model, tokenizer
+
+
+def infer(file, device, model, tokenizer, target_sr):
+    sr, wav = file
+    wav = torch.as_tensor(wav)
+    if wav.dtype == torch.short:
+        wav = wav / 2 ** 15
+    elif wav.dtype == torch.int:
+        wav = wav / 2 ** 31
+    if wav.ndim > 1:
+        wav = wav.mean(1)
+    wav = resample(wav, sr, target_sr)
+    wav_len = len(wav)
+    wav = wav.float().unsqueeze(0).to(device)
+    input_dict = {
+        "mode": "inference",
+        "wav": wav,
+        "wav_len": [wav_len],
+        "specaug": False,
+        "sample_method": "beam",
+        "beam_size": 3,
+    }
+    with torch.no_grad():
+        output_dict = model(input_dict)
+        seq = output_dict["seq"].cpu().numpy()
+        cap = tokenizer.decode(seq)[0]
+    return cap
+
+# def input_toggle(input_type):
+#     if input_type == "file":
+#         return gr.update(visible=True), gr.update(visible=False)
+#     elif input_type == "mic":
+#         return gr.update(visible=False), gr.update(visible=True)
+
+
+if __name__ == "__main__":
+    
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--share", action="store_true", default=False)
+
+    args = parser.parse_args()
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    exp_dir = Path("./checkpoints/audiocaps")
+    cfg = train_util.load_config(exp_dir / "config.yaml")
+    target_sr = cfg["target_sr"]
+    model, tokenizer = load_model(cfg, exp_dir / "ckpt.pth", device)
+
+    with gr.Blocks() as demo:
+        with gr.Row():
+            with gr.Column():
+                # radio = gr.Radio(
+                #     ["file", "mic"],
+                #     value="file",
+                #     label="Select input type"
+                # )
+                file = gr.Audio(label="Input", visible=True)
+                # mic = gr.Microphone(label="Input", visible=False)
+                # radio.change(fn=input_toggle, inputs=radio, outputs=[file, mic])
+                btn = gr.Button("Run")
+            with gr.Column():
+                output = gr.Textbox(label="Output")
+            btn.click(
+                fn=partial(infer,
+                           device=device,
+                           model=model,
+                           tokenizer=tokenizer,
+                           target_sr=target_sr),
+                inputs=[file,],
+                outputs=output
+            )
+        
+        demo.launch(share=args.share)
+

checkpoints/audiocaps/ckpt.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e1c435b1cf05a2b0058dae6f096c4eb4e71c685a19754ed84ea1ee812257434b
+size 55293225

checkpoints/audiocaps/config.yaml

@@ -0,0 +1,30 @@
+tokenizer:
+    type: text_tokenizer.DictTokenizer
+    args:
+        max_length: 20
+
+target_sr: 16000
+
+model:
+    args:
+        shared_dim: 1024
+        tchr_dim: 768
+    model:
+        args: {}
+        decoder:
+            args:
+                attn_emb_dim: 1408
+                dropout: 0.2
+                emb_dim: 256
+                fc_emb_dim: 1408
+                nlayers: 2
+                tie_weights: true
+                vocab_size: 4981
+            type: models.transformer_decoder.TransformerDecoder
+        encoder:
+            args:
+                freeze: false
+                pretrained: true
+            type: models.cnn_encoder.EfficientNetB2
+        type: models.transformer_model.TransformerModel
+    type: models.kd_wrapper.ContraEncoderKdWrapper

models/__init__.py

@@ -0,0 +1,92 @@
+import numpy as np
+import torch
+import torch.nn as nn
+
+from utils.model_util import max_with_lens, mean_with_lens
+
+
+def embedding_pooling(x, lens, pooling="mean"):
+    if pooling == "max":
+        fc_embs = max_with_lens(x, lens)
+    elif pooling == "mean":
+        fc_embs = mean_with_lens(x, lens)
+    elif pooling == "mean+max":
+        x_mean = mean_with_lens(x, lens)
+        x_max = max_with_lens(x, lens)
+        fc_embs = x_mean + x_max
+    elif pooling == "last":
+        indices = (lens - 1).reshape(-1, 1, 1).repeat(1, 1, x.size(-1))
+        # indices: [N, 1, hidden]
+        fc_embs = torch.gather(x, 1, indices).squeeze(1)
+    else:
+        raise Exception(f"pooling method {pooling} not support")
+    return fc_embs
+
+
+class BaseEncoder(nn.Module):
+    
+    """
+    Encode the given audio into embedding
+    Base encoder class, cannot be called directly
+    All encoders should inherit from this class
+    """
+
+    def __init__(self, spec_dim, fc_feat_dim, attn_feat_dim):
+        super(BaseEncoder, self).__init__()
+        self.spec_dim = spec_dim
+        self.fc_feat_dim = fc_feat_dim
+        self.attn_feat_dim = attn_feat_dim
+
+
+    def forward(self, x):
+        #########################
+        # Arguments:
+        # `x`: {
+        #     (may contain)
+        #     wav: [batch_size, n_samples],
+        #     spec: [batch_size, n_frames, spec_dim],
+        #     fc: [batch_size, fc_feat_dim],
+        #     attn: [batch_size, attn_max_len, attn_feat_dim],
+        #     attn_len: [batch_size,]
+        #     ......
+        #  }
+        #
+        # Returns:
+        # `encoded`: {
+        #     fc_emb: [batch_size, fc_emb_dim],
+        #     attn_emb: [batch_size, attn_max_len, attn_emb_dim],
+        #     attn_emb_lens: [batch_size,]
+        # }
+        #########################
+        raise NotImplementedError
+
+
+class BaseDecoder(nn.Module):
+    """
+    Take word/audio embeddings and output the next word probs
+    """
+    def __init__(self, emb_dim, vocab_size, fc_emb_dim,
+                 attn_emb_dim, dropout=0.2, tie_weights=False):
+        super().__init__()
+        self.emb_dim = emb_dim
+        self.vocab_size = vocab_size
+        self.fc_emb_dim = fc_emb_dim
+        self.attn_emb_dim = attn_emb_dim
+        self.tie_weights = tie_weights
+        self.word_embedding = nn.Embedding(vocab_size, emb_dim)
+        self.in_dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        raise NotImplementedError
+
+    def load_word_embedding(self, weight, freeze=True):
+        embedding = np.load(weight)
+        assert embedding.shape[0] == self.vocab_size, "vocabulary size mismatch"
+        assert embedding.shape[1] == self.emb_dim, "embed size mismatch"
+        
+        # embeddings = torch.as_tensor(embeddings).float()
+        # self.word_embeddings.weight = nn.Parameter(embeddings)
+        # for para in self.word_embeddings.parameters():
+            # para.requires_grad = tune
+        self.word_embedding = nn.Embedding.from_pretrained(embedding,
+            freeze=freeze)

models/base.py

@@ -0,0 +1,504 @@
+# -*- coding: utf-8 -*-
+
+from typing import Dict
+
+import torch
+import torch.nn as nn
+
+from utils.model_util import mean_with_lens, repeat_tensor
+
+
+class CaptionMetaMixin:
+    pad_idx = 0
+    start_idx = 1
+    end_idx = 2
+    max_length = 20
+
+    @classmethod
+    def set_index(cls, start_idx, end_idx, pad_idx):
+        cls.start_idx = start_idx
+        cls.end_idx = end_idx
+        cls.pad_idx = pad_idx
+
+
+class CaptionModel(nn.Module, CaptionMetaMixin):
+    """
+    Encoder-decoder captioning model.
+    """
+
+    def __init__(self, encoder: nn.Module, decoder: nn.Module, **kwargs):
+        super().__init__()
+        self.encoder = encoder
+        self.decoder = decoder
+        self.vocab_size = decoder.vocab_size
+        self.train_forward_keys = ["cap", "cap_len", "ss_ratio"]
+        self.inference_forward_keys = ["sample_method", "max_length", "temp"]
+        freeze_encoder = kwargs.get("freeze_encoder", False)
+        if freeze_encoder:
+            for param in self.encoder.parameters():
+                param.requires_grad = False
+        self.check_decoder_compatibility()
+
+    def check_decoder_compatibility(self):
+        compatible_decoders = [x.__class__.__name__ for x in self.compatible_decoders]
+        assert isinstance(self.decoder, self.compatible_decoders), \
+            f"{self.decoder.__class__.__name__} is incompatible with " \
+            f"{self.__class__.__name__}, please use decoder in {compatible_decoders} "
+
+    def forward(self, input_dict: Dict):
+        """
+        input_dict: {
+            (required)
+            mode: train/inference,
+            [spec, spec_len],
+            [fc],
+            [attn, attn_len],
+            [wav, wav_len],
+            [sample_method: greedy],
+            [temp: 1.0] (in case of no teacher forcing)
+
+            (optional, mode=train)
+            cap,
+            cap_len,
+            ss_ratio,
+
+            (optional, mode=inference)
+            sample_method: greedy/beam,
+            max_length,
+            temp,
+            beam_size (optional, sample_method=beam),
+            n_best (optional, sample_method=beam),
+        }
+        """
+        encoder_output_dict = self.encoder(input_dict)
+        output = self.forward_decoder(input_dict, encoder_output_dict)
+        return output
+
+    def forward_decoder(self, input_dict: Dict, encoder_output_dict: Dict):
+        if input_dict["mode"] == "train":
+            forward_dict = {
+                "mode": "train", "sample_method": "greedy", "temp": 1.0
+            }
+            for key in self.train_forward_keys:
+                forward_dict[key] = input_dict[key]
+            forward_dict.update(encoder_output_dict)
+            output = self.train_forward(forward_dict)
+        elif input_dict["mode"] == "inference":
+            forward_dict = {"mode": "inference"}
+            default_args = { "sample_method": "greedy", "max_length": self.max_length, "temp": 1.0 }
+            for key in self.inference_forward_keys:
+                if key in input_dict:
+                    forward_dict[key] = input_dict[key]
+                else:
+                    forward_dict[key] = default_args[key]
+
+            if forward_dict["sample_method"] == "beam":
+                forward_dict["beam_size"] = input_dict.get("beam_size", 3)
+                forward_dict["n_best"] = input_dict.get("n_best", False)
+                forward_dict["n_best_size"] = input_dict.get("n_best_size", forward_dict["beam_size"])
+            elif forward_dict["sample_method"] == "dbs":
+                forward_dict["beam_size"] = input_dict.get("beam_size", 6)
+                forward_dict["group_size"] = input_dict.get("group_size", 3)
+                forward_dict["diversity_lambda"] = input_dict.get("diversity_lambda", 0.5)
+                forward_dict["group_nbest"] = input_dict.get("group_nbest", True)
+
+            forward_dict.update(encoder_output_dict)
+            output = self.inference_forward(forward_dict)
+        else:
+            raise Exception("mode should be either 'train' or 'inference'")
+        output.update(encoder_output_dict)
+        return output
+
+    def prepare_output(self, input_dict):
+        output = {}
+        batch_size = input_dict["fc_emb"].size(0)
+        if input_dict["mode"] == "train":
+            max_length = input_dict["cap"].size(1) - 1
+        elif input_dict["mode"] == "inference":
+            max_length = input_dict["max_length"]
+        else:
+            raise Exception("mode should be either 'train' or 'inference'")
+        device = input_dict["fc_emb"].device
+        output["seq"] = torch.full((batch_size, max_length), self.end_idx,
+                                   dtype=torch.long)
+        output["logit"] = torch.empty(batch_size, max_length,
+                                      self.vocab_size).to(device)
+        output["sampled_logprob"] = torch.zeros(batch_size, max_length)
+        output["embed"] = torch.empty(batch_size, max_length,
+                                      self.decoder.d_model).to(device)
+        return output
+
+    def train_forward(self, input_dict):
+        if input_dict["ss_ratio"] != 1: # scheduled sampling training
+            input_dict["mode"] = "train"
+            return self.stepwise_forward(input_dict)
+        output = self.seq_forward(input_dict)
+        self.train_process(output, input_dict)
+        return output
+
+    def seq_forward(self, input_dict):
+        raise NotImplementedError
+
+    def train_process(self, output, input_dict):
+        pass
+
+    def inference_forward(self, input_dict):
+        if input_dict["sample_method"] == "beam":
+            return self.beam_search(input_dict)
+        elif input_dict["sample_method"] == "dbs":
+            return self.diverse_beam_search(input_dict)
+        return self.stepwise_forward(input_dict)
+
+    def stepwise_forward(self, input_dict):
+        """Step-by-step decoding"""
+        output = self.prepare_output(input_dict)
+        max_length = output["seq"].size(1)
+        # start sampling
+        for t in range(max_length):
+            input_dict["t"] = t
+            self.decode_step(input_dict, output)
+            if input_dict["mode"] == "inference": # decide whether to stop when sampling
+                unfinished_t = output["seq"][:, t] != self.end_idx
+                if t == 0:
+                    unfinished = unfinished_t
+                else:
+                    unfinished *= unfinished_t
+                output["seq"][:, t][~unfinished] = self.end_idx
+                if unfinished.sum() == 0:
+                    break
+        self.stepwise_process(output)
+        return output
+
+    def decode_step(self, input_dict, output):
+        """Decoding operation of timestep t"""
+        decoder_input = self.prepare_decoder_input(input_dict, output)
+        # feed to the decoder to get logit
+        output_t = self.decoder(decoder_input)
+        logit_t = output_t["logit"]
+        # assert logit_t.ndim == 3
+        if logit_t.size(1) == 1:
+            logit_t = logit_t.squeeze(1)
+            embed_t = output_t["embed"].squeeze(1)
+        elif logit_t.size(1) > 1:
+            logit_t = logit_t[:, -1, :]
+            embed_t = output_t["embed"][:, -1, :]
+        else:
+            raise Exception("no logit output")
+        # sample the next input word and get the corresponding logit
+        sampled = self.sample_next_word(logit_t,
+                                        method=input_dict["sample_method"],
+                                        temp=input_dict["temp"])
+
+        output_t.update(sampled)
+        output_t["t"] = input_dict["t"]
+        output_t["logit"] = logit_t
+        output_t["embed"] = embed_t
+        self.stepwise_process_step(output, output_t)
+
+    def prepare_decoder_input(self, input_dict, output):
+        """Prepare the inp ut dict for the decoder"""
+        raise NotImplementedError
+    
+    def stepwise_process_step(self, output, output_t):
+        """Postprocessing (save output values) after each timestep t"""
+        t = output_t["t"]
+        output["logit"][:, t, :] = output_t["logit"]
+        output["seq"][:, t] = output_t["word"]
+        output["sampled_logprob"][:, t] = output_t["probs"]
+        output["embed"][:, t, :] = output_t["embed"]
+
+    def stepwise_process(self, output):
+        """Postprocessing after the whole step-by-step autoregressive decoding"""
+        pass
+
+    def sample_next_word(self, logit, method, temp):
+        """Sample the next word, given probs output by the decoder"""
+        logprob = torch.log_softmax(logit, dim=1)
+        if method == "greedy":
+            sampled_logprob, word = torch.max(logprob.detach(), 1)
+        elif method == "gumbel":
+            def sample_gumbel(shape, eps=1e-20):
+                U = torch.rand(shape).to(logprob.device)
+                return -torch.log(-torch.log(U + eps) + eps)
+            def gumbel_softmax_sample(logit, temperature):
+                y = logit + sample_gumbel(logit.size())
+                return torch.log_softmax(y / temperature, dim=-1)
+            _logprob = gumbel_softmax_sample(logprob, temp)
+            _, word = torch.max(_logprob.data, 1)
+            sampled_logprob = logprob.gather(1, word.unsqueeze(-1))
+        else:
+            logprob = logprob / temp
+            if method.startswith("top"):
+                top_num = float(method[3:])
+                if 0 < top_num < 1: # top-p sampling
+                    probs = torch.softmax(logit, dim=1)
+                    sorted_probs, sorted_indices = torch.sort(probs, descending=True, dim=1)
+                    _cumsum = sorted_probs.cumsum(1)
+                    mask = _cumsum < top_num
+                    mask = torch.cat([torch.ones_like(mask[:,:1]), mask[:,:-1]], 1)
+                    sorted_probs = sorted_probs * mask.to(sorted_probs)
+                    sorted_probs = sorted_probs / sorted_probs.sum(1, keepdim=True)
+                    logprob.scatter_(1, sorted_indices, sorted_probs.log())
+                else: # top-k sampling
+                    k = int(top_num)
+                    tmp = torch.empty_like(logprob).fill_(float('-inf'))
+                    topk, indices = torch.topk(logprob, k, dim=1)
+                    tmp = tmp.scatter(1, indices, topk)
+                    logprob = tmp
+            word = torch.distributions.Categorical(logits=logprob.detach()).sample()
+            sampled_logprob = logprob.gather(1, word.unsqueeze(-1)).squeeze(1)
+        word = word.detach().long()
+        # sampled_logprob: [N,], word: [N,]
+        return {"word": word, "probs": sampled_logprob}
+
+    def beam_search(self, input_dict):
+        output = self.prepare_output(input_dict)
+        max_length = input_dict["max_length"]
+        beam_size = input_dict["beam_size"]
+        if input_dict["n_best"]:
+            n_best_size = input_dict["n_best_size"]
+            batch_size, max_length = output["seq"].size()
+            output["seq"] = torch.full((batch_size, n_best_size, max_length),
+                                        self.end_idx, dtype=torch.long)
+            
+        temp = input_dict["temp"]
+        # instance by instance beam seach
+        for i in range(output["seq"].size(0)):
+            output_i = self.prepare_beamsearch_output(input_dict)
+            input_dict["sample_idx"] = i
+            for t in range(max_length):
+                input_dict["t"] = t
+                output_t = self.beamsearch_step(input_dict, output_i)
+                #######################################
+                # merge with previous beam and select the current max prob beam
+                #######################################
+                logit_t = output_t["logit"]
+                if logit_t.size(1) == 1:
+                    logit_t = logit_t.squeeze(1)
+                elif logit_t.size(1) > 1:
+                    logit_t = logit_t[:, -1, :]
+                else:
+                    raise Exception("no logit output")
+                logprob_t = torch.log_softmax(logit_t, dim=1)
+                logprob_t = torch.log_softmax(logprob_t / temp, dim=1)
+                logprob_t = output_i["topk_logprob"].unsqueeze(1) + logprob_t
+                if t == 0: # for the first step, all k seq will have the same probs
+                    topk_logprob, topk_words = logprob_t[0].topk(
+                        beam_size, 0, True, True)
+                else: # unroll and find top logprob, and their unrolled indices
+                    topk_logprob, topk_words = logprob_t.view(-1).topk(
+                        beam_size, 0, True, True)
+                topk_words = topk_words.cpu()
+                output_i["topk_logprob"] = topk_logprob
+                # output_i["prev_words_beam"] = topk_words // self.vocab_size  # [beam_size,]
+                output_i["prev_words_beam"] = torch.div(topk_words, self.vocab_size,
+                                                        rounding_mode='trunc')
+                output_i["next_word"] = topk_words % self.vocab_size  # [beam_size,]
+                if t == 0:
+                    output_i["seq"] = output_i["next_word"].unsqueeze(1)
+                else:
+                    output_i["seq"] = torch.cat([
+                        output_i["seq"][output_i["prev_words_beam"]],
+                        output_i["next_word"].unsqueeze(1)], dim=1)
+
+                # add finished beams to results
+                is_end = output_i["next_word"] == self.end_idx
+                if t == max_length - 1:
+                    is_end.fill_(1)
+                
+                for beam_idx in range(beam_size):
+                    if is_end[beam_idx]:
+                        final_beam = {
+                            "seq": output_i["seq"][beam_idx].clone(),
+                            "score": output_i["topk_logprob"][beam_idx].item()
+                        }
+                        final_beam["score"] = final_beam["score"] / (t + 1)
+                        output_i["done_beams"].append(final_beam)
+                output_i["topk_logprob"][is_end] -= 1000
+
+                self.beamsearch_process_step(output_i, output_t)
+
+            self.beamsearch_process(output, output_i, input_dict)
+        return output
+
+    def prepare_beamsearch_output(self, input_dict):
+        beam_size = input_dict["beam_size"]
+        device = input_dict["fc_emb"].device
+        output = {
+            "topk_logprob": torch.zeros(beam_size).to(device),
+            "seq": None,
+            "prev_words_beam": None,
+            "next_word": None,
+            "done_beams": [],
+        }
+        return output
+
+    def beamsearch_step(self, input_dict, output_i):
+        decoder_input = self.prepare_beamsearch_decoder_input(input_dict, output_i)
+        output_t = self.decoder(decoder_input)
+        output_t["t"] = input_dict["t"]
+        return output_t
+
+    def prepare_beamsearch_decoder_input(self, input_dict, output_i):
+        raise NotImplementedError
+            
+    def beamsearch_process_step(self, output_i, output_t):
+        pass
+
+    def beamsearch_process(self, output, output_i, input_dict):
+        i = input_dict["sample_idx"]
+        done_beams = sorted(output_i["done_beams"], key=lambda x: -x["score"])
+        if input_dict["n_best"]:
+            done_beams = done_beams[:input_dict["n_best_size"]]
+            for out_idx, done_beam in enumerate(done_beams):
+                seq = done_beam["seq"]
+                output["seq"][i][out_idx, :len(seq)] = seq
+        else:
+            seq = done_beams[0]["seq"]
+            output["seq"][i][:len(seq)] = seq
+    
+    def diverse_beam_search(self, input_dict):
+        
+        def add_diversity(seq_table, logprob, t, divm, diversity_lambda, bdash):
+            local_time = t - divm
+            unaug_logprob = logprob.clone()
+
+            if divm > 0:
+                change = torch.zeros(logprob.size(-1))
+                for prev_choice in range(divm):
+                    prev_decisions = seq_table[prev_choice][..., local_time]
+                    for prev_labels in range(bdash):
+                        change.scatter_add_(0, prev_decisions[prev_labels], change.new_ones(1))
+
+                change = change.to(logprob.device)
+                logprob = logprob - repeat_tensor(change, bdash) * diversity_lambda
+
+            return logprob, unaug_logprob
+
+        output = self.prepare_output(input_dict)
+        group_size = input_dict["group_size"]
+        batch_size = output["seq"].size(0)
+        beam_size = input_dict["beam_size"]
+        bdash = beam_size // group_size
+        input_dict["bdash"] = bdash
+        diversity_lambda = input_dict["diversity_lambda"]
+        device = input_dict["fc_emb"].device
+        max_length = input_dict["max_length"]
+        temp = input_dict["temp"]
+        group_nbest = input_dict["group_nbest"]
+        batch_size, max_length = output["seq"].size()
+        if group_nbest:
+            output["seq"] = torch.full((batch_size, beam_size, max_length),
+                                        self.end_idx, dtype=torch.long)
+        else:
+            output["seq"] = torch.full((batch_size, group_size, max_length),
+                                        self.end_idx, dtype=torch.long)
+
+
+        for i in range(batch_size):
+            input_dict["sample_idx"] = i
+            seq_table = [torch.LongTensor(bdash, 0) for _ in range(group_size)] # group_size x [bdash, 0]
+            logprob_table = [torch.zeros(bdash).to(device) for _ in range(group_size)]
+            done_beams_table = [[] for _ in range(group_size)]
+
+            output_i = {
+                "prev_words_beam": [None for _ in range(group_size)],
+                "next_word": [None for _ in range(group_size)],
+                "state": [None for _ in range(group_size)]
+            }
+
+            for t in range(max_length + group_size - 1):
+                input_dict["t"] = t
+                for divm in range(group_size):
+                    input_dict["divm"] = divm
+                    if t >= divm and t <= max_length + divm - 1:
+                        local_time = t - divm
+                        decoder_input = self.prepare_dbs_decoder_input(input_dict, output_i)
+                        output_t = self.decoder(decoder_input)
+                        output_t["divm"] = divm
+                        logit_t = output_t["logit"]
+                        if logit_t.size(1) == 1:
+                            logit_t = logit_t.squeeze(1)
+                        elif logit_t.size(1) > 1:
+                            logit_t = logit_t[:, -1, :]
+                        else:
+                            raise Exception("no logit output")
+                        logprob_t = torch.log_softmax(logit_t, dim=1)
+                        logprob_t = torch.log_softmax(logprob_t / temp, dim=1)
+                        logprob_t, unaug_logprob_t = add_diversity(seq_table, logprob_t, t, divm, diversity_lambda, bdash)
+                        logprob_t = logprob_table[divm].unsqueeze(-1) + logprob_t
+                        if local_time == 0: # for the first step, all k seq will have the same probs
+                            topk_logprob, topk_words = logprob_t[0].topk(
+                                bdash, 0, True, True)
+                        else: # unroll and find top logprob, and their unrolled indices
+                            topk_logprob, topk_words = logprob_t.view(-1).topk(
+                                bdash, 0, True, True)
+                        topk_words = topk_words.cpu()
+                        logprob_table[divm] = topk_logprob
+                        output_i["prev_words_beam"][divm] = topk_words // self.vocab_size  # [bdash,]
+                        output_i["next_word"][divm] = topk_words % self.vocab_size  # [bdash,]
+                        if local_time > 0:
+                            seq_table[divm] = seq_table[divm][output_i["prev_words_beam"][divm]]
+                        seq_table[divm] = torch.cat([
+                            seq_table[divm],
+                            output_i["next_word"][divm].unsqueeze(-1)], -1)
+
+                        is_end = seq_table[divm][:, t-divm] == self.end_idx
+                        assert seq_table[divm].shape[-1] == t - divm + 1
+                        if t == max_length + divm - 1:
+                            is_end.fill_(1)
+                        for beam_idx in range(bdash):
+                            if is_end[beam_idx]:
+                                final_beam = {
+                                    "seq": seq_table[divm][beam_idx].clone(),
+                                    "score": logprob_table[divm][beam_idx].item()
+                                }
+                                final_beam["score"] = final_beam["score"] / (t - divm + 1)
+                                done_beams_table[divm].append(final_beam)
+                        logprob_table[divm][is_end] -= 1000
+                        self.dbs_process_step(output_i, output_t)
+            done_beams_table = [sorted(done_beams_table[divm], key=lambda x: -x["score"])[:bdash] for divm in range(group_size)]
+            if group_nbest:
+                done_beams = sum(done_beams_table, [])
+            else:
+                done_beams = [group_beam[0] for group_beam in done_beams_table]
+            for _, done_beam in enumerate(done_beams):
+                output["seq"][i, _, :len(done_beam["seq"])] = done_beam["seq"]
+
+        return output
+            
+    def prepare_dbs_decoder_input(self, input_dict, output_i):
+        raise NotImplementedError
+
+    def dbs_process_step(self, output_i, output_t):
+        pass
+
+
+class CaptionSequenceModel(nn.Module, CaptionMetaMixin):
+
+    def __init__(self, model, seq_output_size):
+        super().__init__()
+        self.model = model
+        if model.decoder.d_model != seq_output_size:
+            self.output_transform = nn.Linear(model.decoder.d_model, seq_output_size)
+        else:
+            self.output_transform = lambda x: x
+
+    def forward(self, input_dict):
+        output = self.model(input_dict)
+
+        if input_dict["mode"] == "train":
+            lens = input_dict["cap_len"] - 1
+            # seq_outputs: [N, d_model]
+        elif input_dict["mode"] == "inference":
+            if "sample_method" in input_dict and input_dict["sample_method"] == "beam":
+                return output
+            seq = output["seq"]
+            lens = torch.where(seq == self.model.end_idx, torch.zeros_like(seq), torch.ones_like(seq)).sum(dim=1)
+        else:
+            raise Exception("mode should be either 'train' or 'inference'")
+        seq_output = mean_with_lens(output["embed"], lens)
+        seq_output = self.output_transform(seq_output)
+        output["seq_output"] = seq_output
+        return output
+

models/cnn_encoder.py

@@ -0,0 +1,810 @@
+# -*- coding: utf-8 -*-
+
+from einops import rearrange
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchaudio import transforms
+
+from utils.model_util import mean_with_lens, max_with_lens
+from utils.train_util import merge_load_state_dict
+
+
+def init_layer(layer):
+    """Initialize a Linear or Convolutional layer. """
+    nn.init.xavier_uniform_(layer.weight)
+ 
+    if hasattr(layer, 'bias'):
+        if layer.bias is not None:
+            layer.bias.data.fill_(0.)
+            
+    
+def init_bn(bn):
+    """Initialize a Batchnorm layer. """
+    bn.bias.data.fill_(0.)
+    bn.weight.data.fill_(1.)
+
+
+class ConvBlock(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        
+        super(ConvBlock, self).__init__()
+        
+        self.conv1 = nn.Conv2d(in_channels=in_channels,
+                              out_channels=out_channels,
+                              kernel_size=(3, 3), stride=(1, 1),
+                              padding=(1, 1), bias=False)
+                              
+        self.conv2 = nn.Conv2d(in_channels=out_channels,
+                              out_channels=out_channels,
+                              kernel_size=(3, 3), stride=(1, 1),
+                              padding=(1, 1), bias=False)
+                              
+        self.bn1 = nn.BatchNorm2d(out_channels)
+        self.bn2 = nn.BatchNorm2d(out_channels)
+
+        self.init_weight()
+        
+    def init_weight(self):
+        init_layer(self.conv1)
+        init_layer(self.conv2)
+        init_bn(self.bn1)
+        init_bn(self.bn2)
+
+        
+    def forward(self, input, pool_size=(2, 2), pool_type='avg'):
+        
+        x = input
+        x = F.relu_(self.bn1(self.conv1(x)))
+        x = F.relu_(self.bn2(self.conv2(x)))
+        if pool_type == 'max':
+            x = F.max_pool2d(x, kernel_size=pool_size)
+        elif pool_type == 'avg':
+            x = F.avg_pool2d(x, kernel_size=pool_size)
+        elif pool_type == 'avg+max':
+            x1 = F.avg_pool2d(x, kernel_size=pool_size)
+            x2 = F.max_pool2d(x, kernel_size=pool_size)
+            x = x1 + x2
+        else:
+            raise Exception('Incorrect argument!')
+        
+        return x
+
+
+class ConvBlock5x5(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        
+        super(ConvBlock5x5, self).__init__()
+        
+        self.conv1 = nn.Conv2d(in_channels=in_channels, 
+                              out_channels=out_channels,
+                              kernel_size=(5, 5), stride=(1, 1),
+                              padding=(2, 2), bias=False)
+                              
+        self.bn1 = nn.BatchNorm2d(out_channels)
+
+        self.init_weight()
+        
+    def init_weight(self):
+        init_layer(self.conv1)
+        init_bn(self.bn1)
+
+    def forward(self, input, pool_size=(2, 2), pool_type='avg'):
+        
+        x = input
+        x = F.relu_(self.bn1(self.conv1(x)))
+        if pool_type == 'max':
+            x = F.max_pool2d(x, kernel_size=pool_size)
+        elif pool_type == 'avg':
+            x = F.avg_pool2d(x, kernel_size=pool_size)
+        elif pool_type == 'avg+max':
+            x1 = F.avg_pool2d(x, kernel_size=pool_size)
+            x2 = F.max_pool2d(x, kernel_size=pool_size)
+            x = x1 + x2
+        else:
+            raise Exception('Incorrect argument!')
+        
+        return x
+
+
+class Cnn6Encoder(nn.Module):
+
+    def __init__(self, sample_rate=32000, freeze=False):
+        super().__init__()
+
+        sr_to_fmax = {
+            32000: 14000,
+            16000: 8000
+        }
+        # Logmel spectrogram extractor
+        self.melspec_extractor = transforms.MelSpectrogram(
+            sample_rate=sample_rate,
+            n_fft=32 * sample_rate // 1000,
+            win_length=32 * sample_rate // 1000,
+            hop_length=10 * sample_rate // 1000,
+            f_min=50,
+            f_max=sr_to_fmax[sample_rate],
+            n_mels=64,
+            norm="slaney",
+            mel_scale="slaney"
+        )
+        self.hop_length = 10 * sample_rate // 1000
+        self.db_transform = transforms.AmplitudeToDB()
+
+        self.bn0 = nn.BatchNorm2d(64)
+
+        self.conv_block1 = ConvBlock5x5(in_channels=1, out_channels=64)
+        self.conv_block2 = ConvBlock5x5(in_channels=64, out_channels=128)
+        self.conv_block3 = ConvBlock5x5(in_channels=128, out_channels=256)
+        self.conv_block4 = ConvBlock5x5(in_channels=256, out_channels=512)
+
+        self.downsample_ratio = 16
+
+        self.fc1 = nn.Linear(512, 512, bias=True)
+        self.fc_emb_size = 512
+        self.init_weight()
+        self.freeze = freeze
+
+    def init_weight(self):
+        init_bn(self.bn0)
+        init_layer(self.fc1)
+
+    def load_pretrained(self, pretrained, output_fn):
+        checkpoint = torch.load(pretrained, map_location="cpu")
+
+        if "model" in checkpoint:
+            state_dict = checkpoint["model"]
+        else:
+            raise Exception("Unkown checkpoint format")
+
+        loaded_keys = merge_load_state_dict(state_dict, self, output_fn)
+        if self.freeze:
+            for name, param in self.named_parameters():
+                if name in loaded_keys:
+                    param.requires_grad = False
+                else:
+                    param.requires_grad = True
+
+    def forward(self, input_dict):
+        waveform = input_dict["wav"]
+        wave_length = input_dict["wav_len"]
+        specaug = input_dict["specaug"]
+        x = self.melspec_extractor(waveform)
+        x = self.db_transform(x)    # (batch_size, mel_bins, time_steps)
+        x = x.transpose(1, 2)
+        x = x.unsqueeze(1)      # (batch_size, 1, time_steps, mel_bins)
+
+        x = x.transpose(1, 3)
+        x = self.bn0(x)
+        x = x.transpose(1, 3)
+        
+        x = self.conv_block1(x, pool_size=(2, 2), pool_type='avg')
+        x = F.dropout(x, p=0.2, training=self.training)
+        x = self.conv_block2(x, pool_size=(2, 2), pool_type='avg')
+        x = F.dropout(x, p=0.2, training=self.training)
+        x = self.conv_block3(x, pool_size=(2, 2), pool_type='avg')
+        x = F.dropout(x, p=0.2, training=self.training)
+        x = self.conv_block4(x, pool_size=(2, 2), pool_type='avg')
+        x = F.dropout(x, p=0.2, training=self.training)
+
+        x = torch.mean(x, dim=3)
+        attn_emb = x.transpose(1, 2)
+        wave_length = torch.as_tensor(wave_length)
+        feat_length = torch.div(wave_length, self.hop_length,
+            rounding_mode="floor") + 1
+        feat_length = torch.div(feat_length, self.downsample_ratio,
+            rounding_mode="floor")
+        x_max = max_with_lens(attn_emb, feat_length)
+        x_mean = mean_with_lens(attn_emb, feat_length)
+        x = x_max + x_mean
+        x = F.dropout(x, p=0.5, training=self.training)
+        x = F.relu_(self.fc1(x))
+        fc_emb = F.dropout(x, p=0.5, training=self.training)
+
+        return {
+            "attn_emb": attn_emb,
+            "fc_emb": fc_emb,
+            "attn_emb_len": feat_length
+        }
+
+
+class Cnn10Encoder(nn.Module):
+
+    def __init__(self, sample_rate=32000, freeze=False):
+        super().__init__()
+
+        sr_to_fmax = {
+            32000: 14000,
+            16000: 8000
+        }
+        # Logmel spectrogram extractor
+        self.melspec_extractor = transforms.MelSpectrogram(
+            sample_rate=sample_rate,
+            n_fft=32 * sample_rate // 1000,
+            win_length=32 * sample_rate // 1000,
+            hop_length=10 * sample_rate // 1000,
+            f_min=50,
+            f_max=sr_to_fmax[sample_rate],
+            n_mels=64,
+            norm="slaney",
+            mel_scale="slaney"
+        )
+        self.hop_length = 10 * sample_rate // 1000
+        self.db_transform = transforms.AmplitudeToDB()
+
+        self.bn0 = nn.BatchNorm2d(64)
+
+        self.conv_block1 = ConvBlock(in_channels=1, out_channels=64)
+        self.conv_block2 = ConvBlock(in_channels=64, out_channels=128)
+        self.conv_block3 = ConvBlock(in_channels=128, out_channels=256)
+        self.conv_block4 = ConvBlock(in_channels=256, out_channels=512)
+
+        self.downsample_ratio = 16
+
+        self.fc1 = nn.Linear(512, 512, bias=True)
+        self.fc_emb_size = 512
+        self.init_weight()
+        self.freeze = freeze
+
+    def init_weight(self):
+        init_bn(self.bn0)
+        init_layer(self.fc1)
+
+    def load_pretrained(self, pretrained, output_fn):
+        checkpoint = torch.load(pretrained, map_location="cpu")
+
+        if "model" in checkpoint:
+            state_dict = checkpoint["model"]
+        else:
+            raise Exception("Unkown checkpoint format")
+
+        loaded_keys = merge_load_state_dict(state_dict, self, output_fn)
+        if self.freeze:
+            for name, param in self.named_parameters():
+                if name in loaded_keys:
+                    param.requires_grad = False
+                else:
+                    param.requires_grad = True
+
+    def forward(self, input_dict):
+        waveform = input_dict["wav"]
+        wave_length = input_dict["wav_len"]
+        specaug = input_dict["specaug"]
+        x = self.melspec_extractor(waveform)
+        x = self.db_transform(x)    # (batch_size, mel_bins, time_steps)
+        x = x.transpose(1, 2)
+        x = x.unsqueeze(1)      # (batch_size, 1, time_steps, mel_bins)
+
+        x = x.transpose(1, 3)
+        x = self.bn0(x)
+        x = x.transpose(1, 3)
+        
+        x = self.conv_block1(x, pool_size=(2, 2), pool_type='avg')
+        x = F.dropout(x, p=0.2, training=self.training)
+        x = self.conv_block2(x, pool_size=(2, 2), pool_type='avg')
+        x = F.dropout(x, p=0.2, training=self.training)
+        x = self.conv_block3(x, pool_size=(2, 2), pool_type='avg')
+        x = F.dropout(x, p=0.2, training=self.training)
+        x = self.conv_block4(x, pool_size=(2, 2), pool_type='avg')
+        x = F.dropout(x, p=0.2, training=self.training)
+
+        x = torch.mean(x, dim=3)
+        attn_emb = x.transpose(1, 2)
+        wave_length = torch.as_tensor(wave_length)
+        feat_length = torch.div(wave_length, self.hop_length,
+            rounding_mode="floor") + 1
+        feat_length = torch.div(feat_length, self.downsample_ratio,
+            rounding_mode="floor")
+        x_max = max_with_lens(attn_emb, feat_length)
+        x_mean = mean_with_lens(attn_emb, feat_length)
+        x = x_max + x_mean
+        x = F.dropout(x, p=0.5, training=self.training)
+        x = F.relu_(self.fc1(x))
+        fc_emb = F.dropout(x, p=0.5, training=self.training)
+
+        return {
+            "attn_emb": attn_emb,
+            "fc_emb": fc_emb,
+            "attn_emb_len": feat_length
+        }
+
+
+class Cnn14Encoder(nn.Module):
+    def __init__(self, sample_rate=32000, freeze=False):
+        super().__init__()
+        sr_to_fmax = {
+            32000: 14000,
+            16000: 8000
+        }
+        # Logmel spectrogram extractor
+        self.melspec_extractor = transforms.MelSpectrogram(
+            sample_rate=sample_rate,
+            n_fft=32 * sample_rate // 1000,
+            win_length=32 * sample_rate // 1000,
+            hop_length=10 * sample_rate // 1000,
+            f_min=50,
+            f_max=sr_to_fmax[sample_rate],
+            n_mels=64,
+            norm="slaney",
+            mel_scale="slaney"
+        )
+        self.hop_length = 10 * sample_rate // 1000
+        self.db_transform = transforms.AmplitudeToDB()
+
+        self.bn0 = nn.BatchNorm2d(64)
+
+        self.conv_block1 = ConvBlock(in_channels=1, out_channels=64)
+        self.conv_block2 = ConvBlock(in_channels=64, out_channels=128)
+        self.conv_block3 = ConvBlock(in_channels=128, out_channels=256)
+        self.conv_block4 = ConvBlock(in_channels=256, out_channels=512)
+        self.conv_block5 = ConvBlock(in_channels=512, out_channels=1024)
+        self.conv_block6 = ConvBlock(in_channels=1024, out_channels=2048)
+
+        self.downsample_ratio = 32
+
+        self.fc1 = nn.Linear(2048, 2048, bias=True)
+        self.fc_emb_size = 2048
+        
+        self.init_weight()
+        self.freeze = freeze
+
+    def init_weight(self):
+        init_bn(self.bn0)
+        init_layer(self.fc1)
+
+    def load_pretrained(self, pretrained, output_fn):
+        checkpoint = torch.load(pretrained, map_location="cpu")
+
+        if "model" in checkpoint:
+            state_keys = checkpoint["model"].keys()
+            backbone = False
+            for key in state_keys:
+                if key.startswith("backbone."):
+                    backbone = True
+                    break
+
+            if backbone: # COLA
+                state_dict = {}
+                for key, value in checkpoint["model"].items():
+                    if key.startswith("backbone."):
+                        model_key = key.replace("backbone.", "")
+                        state_dict[model_key] = value
+            else: # PANNs
+                state_dict = checkpoint["model"]
+        elif "state_dict" in checkpoint: # BLAT
+            state_dict = checkpoint["state_dict"]
+            state_dict_keys = list(filter(
+                lambda x: "audio_encoder" in x, state_dict.keys()))
+            state_dict = {
+                key.replace('audio_encoder.', ''): state_dict[key]
+                    for key in state_dict_keys
+            }
+        else:
+            raise Exception("Unkown checkpoint format")
+
+        loaded_keys = merge_load_state_dict(state_dict, self, output_fn)
+        if self.freeze:
+            for name, param in self.named_parameters():
+                if name in loaded_keys:
+                    param.requires_grad = False
+                else:
+                    param.requires_grad = True
+ 
+    def forward(self, input_dict):
+        waveform = input_dict["wav"]
+        wave_length = input_dict["wav_len"]
+        specaug = input_dict["specaug"]
+        x = self.melspec_extractor(waveform)
+        x = self.db_transform(x)    # (batch_size, mel_bins, time_steps)
+        x = x.transpose(1, 2)
+        x = x.unsqueeze(1)      # (batch_size, 1, time_steps, mel_bins)
+
+        x = x.transpose(1, 3)
+        x = self.bn0(x)
+        x = x.transpose(1, 3)
+
+        x = self.conv_block1(x, pool_size=(2, 2), pool_type='avg')
+        x = F.dropout(x, p=0.2, training=self.training)
+        x = self.conv_block2(x, pool_size=(2, 2), pool_type='avg')
+        x = F.dropout(x, p=0.2, training=self.training)
+        x = self.conv_block3(x, pool_size=(2, 2), pool_type='avg')
+        x = F.dropout(x, p=0.2, training=self.training)
+        x = self.conv_block4(x, pool_size=(2, 2), pool_type='avg')
+        x = F.dropout(x, p=0.2, training=self.training)
+        x = self.conv_block5(x, pool_size=(2, 2), pool_type='avg')
+        x = F.dropout(x, p=0.2, training=self.training)
+        x = self.conv_block6(x, pool_size=(1, 1), pool_type='avg')
+        x = F.dropout(x, p=0.2, training=self.training)
+        x = torch.mean(x, dim=3)
+        attn_emb = x.transpose(1, 2)
+        
+        wave_length = torch.as_tensor(wave_length)
+        feat_length = torch.div(wave_length, self.hop_length,
+            rounding_mode="floor") + 1
+        feat_length = torch.div(feat_length, self.downsample_ratio,
+            rounding_mode="floor")
+        x_max = max_with_lens(attn_emb, feat_length)
+        x_mean = mean_with_lens(attn_emb, feat_length)
+        x = x_max + x_mean
+        x = F.dropout(x, p=0.5, training=self.training)
+        x = F.relu_(self.fc1(x))
+        fc_emb = F.dropout(x, p=0.5, training=self.training)
+        
+        output_dict = {
+            'fc_emb': fc_emb,
+            'attn_emb': attn_emb,
+            'attn_emb_len': feat_length
+        }
+
+        return output_dict
+
+
+class InvertedResidual(nn.Module):
+
+    def __init__(self, inp, oup, stride, expand_ratio):
+        super().__init__()
+        self.stride = stride
+        assert stride in [1, 2]
+
+        hidden_dim = round(inp * expand_ratio)
+        self.use_res_connect = self.stride == 1 and inp == oup
+
+        if expand_ratio == 1:
+            _layers = [
+                nn.Conv2d(hidden_dim, hidden_dim, 3, 1, 1, groups=hidden_dim, bias=False), 
+                nn.AvgPool2d(stride), 
+                nn.BatchNorm2d(hidden_dim), 
+                nn.ReLU6(inplace=True), 
+                nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False), 
+                nn.BatchNorm2d(oup)
+                ]
+            _layers = nn.Sequential(*_layers)
+            init_layer(_layers[0])
+            init_bn(_layers[2])
+            init_layer(_layers[4])
+            init_bn(_layers[5])
+            self.conv = _layers
+        else:
+            _layers = [
+                nn.Conv2d(inp, hidden_dim, 1, 1, 0, bias=False), 
+                nn.BatchNorm2d(hidden_dim), 
+                nn.ReLU6(inplace=True), 
+                nn.Conv2d(hidden_dim, hidden_dim, 3, 1, 1, groups=hidden_dim, bias=False), 
+                nn.AvgPool2d(stride), 
+                nn.BatchNorm2d(hidden_dim), 
+                nn.ReLU6(inplace=True), 
+                nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False), 
+                nn.BatchNorm2d(oup)
+                ]
+            _layers = nn.Sequential(*_layers)
+            init_layer(_layers[0])
+            init_bn(_layers[1])
+            init_layer(_layers[3])
+            init_bn(_layers[5])
+            init_layer(_layers[7])
+            init_bn(_layers[8])
+            self.conv = _layers
+
+    def forward(self, x):
+        if self.use_res_connect:
+            return x + self.conv(x)
+        else:
+            return self.conv(x)
+
+
+class MobileNetV2(nn.Module):
+    def __init__(self, sample_rate):
+        
+        super().__init__()
+
+        sr_to_fmax = {
+            32000: 14000,
+            16000: 8000
+        }
+        # Logmel spectrogram extractor
+        self.melspec_extractor = transforms.MelSpectrogram(
+            sample_rate=sample_rate,
+            n_fft=32 * sample_rate // 1000,
+            win_length=32 * sample_rate // 1000,
+            hop_length=10 * sample_rate // 1000,
+            f_min=50,
+            f_max=sr_to_fmax[sample_rate],
+            n_mels=64,
+            norm="slaney",
+            mel_scale="slaney"
+        )
+        self.hop_length = 10 * sample_rate // 1000
+        self.db_transform = transforms.AmplitudeToDB()
+
+        self.bn0 = nn.BatchNorm2d(64)
+ 
+        width_mult=1.
+        block = InvertedResidual
+        input_channel = 32
+        last_channel = 1280
+        interverted_residual_setting = [
+            # t, c, n, s
+            [1, 16, 1, 1],
+            [6, 24, 2, 2],
+            [6, 32, 3, 2],
+            [6, 64, 4, 2],
+            [6, 96, 3, 2],
+            [6, 160, 3, 1],
+            [6, 320, 1, 1],
+        ]
+
+        self.downsample_ratio = 32
+
+        def conv_bn(inp, oup, stride):
+            _layers = [
+                nn.Conv2d(inp, oup, 3, 1, 1, bias=False), 
+                nn.AvgPool2d(stride), 
+                nn.BatchNorm2d(oup), 
+                nn.ReLU6(inplace=True)
+                ]
+            _layers = nn.Sequential(*_layers)
+            init_layer(_layers[0])
+            init_bn(_layers[2])
+            return _layers
+
+
+        def conv_1x1_bn(inp, oup):
+            _layers = nn.Sequential(
+                nn.Conv2d(inp, oup, 1, 1, 0, bias=False),
+                nn.BatchNorm2d(oup),
+                nn.ReLU6(inplace=True)
+            )
+            init_layer(_layers[0])
+            init_bn(_layers[1])
+            return _layers
+
+        # building first layer
+        input_channel = int(input_channel * width_mult)
+        self.last_channel = int(last_channel * width_mult) if width_mult > 1.0 else last_channel
+        self.features = [conv_bn(1, input_channel, 2)]
+        # building inverted residual blocks
+        for t, c, n, s in interverted_residual_setting:
+            output_channel = int(c * width_mult)
+            for i in range(n):
+                if i == 0:
+                    self.features.append(block(input_channel, output_channel, s, expand_ratio=t))
+                else:
+                    self.features.append(block(input_channel, output_channel, 1, expand_ratio=t))
+                input_channel = output_channel
+        # building last several layers
+        self.features.append(conv_1x1_bn(input_channel, self.last_channel))
+        # make it nn.Sequential
+        self.features = nn.Sequential(*self.features)
+
+        self.fc1 = nn.Linear(1280, 1024, bias=True)
+        
+        self.init_weight()
+
+    def init_weight(self):
+        init_bn(self.bn0)
+        init_layer(self.fc1)
+ 
+    def forward(self, input_dict):
+
+        waveform = input_dict["wav"]
+        wave_length = input_dict["wav_len"]
+        specaug = input_dict["specaug"]
+        x = self.melspec_extractor(waveform)
+        x = self.db_transform(x)    # (batch_size, mel_bins, time_steps)
+        x = x.transpose(1, 2)
+        x = x.unsqueeze(1)      # (batch_size, 1, time_steps, mel_bins)
+        
+        x = x.transpose(1, 3)
+        x = self.bn0(x)
+        x = x.transpose(1, 3)
+        
+        x = self.features(x)
+        
+        x = torch.mean(x, dim=3)
+        attn_emb = x.transpose(1, 2)
+        
+        wave_length = torch.as_tensor(wave_length)
+        feat_length = torch.div(wave_length, self.hop_length,
+            rounding_mode="floor") + 1
+        feat_length = torch.div(feat_length, self.downsample_ratio,
+            rounding_mode="floor")
+        x_max = max_with_lens(attn_emb, feat_length)
+        x_mean = mean_with_lens(attn_emb, feat_length)
+        x = x_max + x_mean
+        # TODO: the original PANNs code does not have dropout here, why?
+        x = F.dropout(x, p=0.5, training=self.training)
+        x = F.relu_(self.fc1(x))
+        fc_emb = F.dropout(x, p=0.5, training=self.training)
+        
+        output_dict = {
+            'fc_emb': fc_emb,
+            'attn_emb': attn_emb,
+            'attn_emb_len': feat_length
+        }
+
+        return output_dict
+
+
+class MobileNetV3(nn.Module):
+    
+    def __init__(self,
+                 sample_rate,
+                 model_name,
+                 n_mels=64,
+                 win_length=32,
+                 pretrained=True,
+                 freeze=False,
+                 pooling="mean_max_fc"):
+
+        from captioning.models.eff_at_encoder import get_model, NAME_TO_WIDTH
+
+        super().__init__()
+        sr_to_fmax = {
+            32000: 14000,
+            16000: 8000
+        }
+        self.n_mels = n_mels
+        # Logmel spectrogram extractor
+        self.melspec_extractor = transforms.MelSpectrogram(
+            sample_rate=sample_rate,
+            n_fft=32 * sample_rate // 1000,
+            win_length=win_length * sample_rate // 1000,
+            hop_length=10 * sample_rate // 1000,
+            f_min=50,
+            f_max=sr_to_fmax[sample_rate],
+            n_mels=n_mels,
+            norm="slaney",
+            mel_scale="slaney"
+        )
+        self.hop_length = 10 * sample_rate // 1000
+        self.db_transform = transforms.AmplitudeToDB()
+
+        self.bn0 = nn.BatchNorm2d(n_mels)
+        
+        width_mult = NAME_TO_WIDTH(model_name)
+        self.features = get_model(model_name=model_name,
+                                  pretrained=pretrained,
+                                  width_mult=width_mult).features
+        self.downsample_ratio = 32
+
+        if pooling == "mean_max_fc":
+            self.fc_emb_size = 512
+            self.fc1 = nn.Linear(self.features[-1].out_channels, 512, bias=True)
+        elif pooling == "mean":
+            self.fc_emb_size = self.features[-1].out_channels
+        self.init_weight()
+
+        if freeze:
+            for param in self.parameters():
+                param.requires_grad = False
+
+        self.pooling = pooling
+
+    def init_weight(self):
+        init_bn(self.bn0)
+        if hasattr(self, "fc1"):
+            init_layer(self.fc1)
+
+    def forward(self, input_dict):
+
+        waveform = input_dict["wav"]
+        wave_length = input_dict["wav_len"]
+        specaug = input_dict["specaug"]
+        x = self.melspec_extractor(waveform)
+        x = self.db_transform(x)    # (batch_size, mel_bins, time_steps)
+        x = x.transpose(1, 2)
+        x = x.unsqueeze(1)      # (batch_size, 1, time_steps, mel_bins)
+        
+        x = x.transpose(1, 3)
+        x = self.bn0(x)
+        x = x.transpose(1, 3)
+        
+        x = self.features(x)
+        
+        x = torch.mean(x, dim=3)
+        attn_emb = x.transpose(1, 2)
+        
+        wave_length = torch.as_tensor(wave_length)
+        feat_length = torch.div(wave_length, self.hop_length,
+            rounding_mode="floor") + 1
+        feat_length = torch.div(feat_length, self.downsample_ratio,
+            rounding_mode="floor")
+
+        if self.pooling == "mean_max_fc":
+            x_max = max_with_lens(attn_emb, feat_length)
+            x_mean = mean_with_lens(attn_emb, feat_length)
+            x = x_max + x_mean
+            x = F.dropout(x, p=0.5, training=self.training)
+            x = F.relu_(self.fc1(x))
+            fc_emb = F.dropout(x, p=0.5, training=self.training)
+        elif self.pooling == "mean":
+            fc_emb = mean_with_lens(attn_emb, feat_length)
+        
+        output_dict = {
+            'fc_emb': fc_emb,
+            'attn_emb': attn_emb,
+            'attn_emb_len': feat_length
+        }
+
+        return output_dict
+
+
+class EfficientNetB2(nn.Module):
+
+    def __init__(self,
+                 n_mels: int = 64,
+                 win_length: int = 32,
+                 hop_length: int = 10,
+                 f_min: int = 0,
+                 pretrained: bool = False,
+                 prune_ratio: float = 0.0,
+                 prune_se: bool = True,
+                 prune_start_layer: int = 0,
+                 prune_method: str = "operator_norm",
+                 freeze: bool = False,):
+        from models.eff_latent_encoder import get_model, get_pruned_model
+        super().__init__()
+        sample_rate = 16000
+        self.melspec_extractor = transforms.MelSpectrogram(
+            sample_rate=sample_rate,
+            n_fft=win_length * sample_rate // 1000,
+            win_length=win_length * sample_rate // 1000,
+            hop_length=hop_length * sample_rate // 1000,
+            f_min=f_min,
+            n_mels=n_mels,
+        )
+        self.hop_length = 10 * sample_rate // 1000
+        self.db_transform = transforms.AmplitudeToDB(top_db=120)
+        if prune_ratio > 0:
+            self.backbone = get_pruned_model(pretrained=pretrained,
+                                             prune_ratio=prune_ratio,
+                                             prune_start_layer=prune_start_layer,
+                                             prune_se=prune_se,
+                                             prune_method=prune_method)
+        else:
+            self.backbone = get_model(pretrained=pretrained)
+        self.fc_emb_size = self.backbone.eff_net._conv_head.out_channels
+        self.downsample_ratio = 32
+        if freeze:
+            for param in self.parameters():
+                param.requires_grad = False
+
+    def forward(self, input_dict):
+        
+        waveform = input_dict["wav"]
+        wave_length = input_dict["wav_len"]
+        specaug = input_dict["specaug"]
+        x = self.melspec_extractor(waveform)
+        x = self.db_transform(x)    # (batch_size, mel_bins, time_steps)
+        
+        x = self.backbone(x)
+        attn_emb = x
+
+        wave_length = torch.as_tensor(wave_length)
+        feat_length = torch.div(wave_length, self.hop_length,
+            rounding_mode="floor") + 1
+        feat_length = torch.div(feat_length, self.downsample_ratio,
+            rounding_mode="floor")
+        fc_emb = mean_with_lens(attn_emb, feat_length)
+        
+        output_dict = {
+            'fc_emb': fc_emb,
+            'attn_emb': attn_emb,
+            'attn_emb_len': feat_length
+        }
+        return output_dict
+
+
+if __name__ == "__main__":
+    encoder = MobileNetV3(32000, "mn10_as")
+    print(encoder)
+    input_dict = {
+        "wav": torch.randn(4, 320000), 
+        "wav_len": torch.tensor([320000, 280000, 160000, 300000]),
+        "specaug": True
+    }
+    output_dict = encoder(input_dict)
+    print("attn embed: ", output_dict["attn_emb"].shape)
+    print("fc embed: ", output_dict["fc_emb"].shape)
+    print("attn embed length: ", output_dict["attn_emb_len"])

models/eff_latent_encoder.py

@@ -0,0 +1,347 @@
+import os
+
+import torch
+import torch.nn as nn
+from tqdm import tqdm
+from efficientnet_pytorch import EfficientNet
+from efficientnet_pytorch.model import MBConvBlock
+from efficientnet_pytorch import utils as efficientnet_utils
+from efficientnet_pytorch.utils import (
+    round_filters,
+    round_repeats,
+    get_same_padding_conv2d,
+    calculate_output_image_size,
+    MemoryEfficientSwish,
+)
+from einops import rearrange, reduce
+from torch.hub import load_state_dict_from_url
+
+
+model_dir = "./"
+
+
+class _EffiNet(nn.Module):
+    """A proxy for efficient net models"""
+    def __init__(self,
+                 blocks_args=None,
+                 global_params=None,
+                 prune_start_layer: int = 0,
+                 prune_se: bool = True,
+                 prune_ratio: float = 0.0
+                 ) -> None:
+        super().__init__()
+        if prune_ratio > 0:
+            self.eff_net = EfficientNetB2Pruned(blocks_args=blocks_args,
+                                                global_params=global_params,
+                                                prune_start_layer=prune_start_layer,
+                                                prune_se=prune_se,
+                                                prune_ratio=prune_ratio)
+        else:
+            self.eff_net = EfficientNet(blocks_args=blocks_args,
+                                        global_params=global_params)
+        
+
+    def forward(self, x: torch.Tensor): 
+        x = rearrange(x, 'b f t -> b 1 f t')
+        x = self.eff_net.extract_features(x)
+        return reduce(x, 'b c f t -> b t c', 'mean')
+
+
+def get_model(pretrained=True) -> _EffiNet:
+    blocks_args, global_params = efficientnet_utils.get_model_params(
+        'efficientnet-b2', {'include_top': False})
+    model = _EffiNet(blocks_args=blocks_args,
+                     global_params=global_params)
+    model.eff_net._change_in_channels(1)
+    if pretrained:
+        model_path = os.path.join(model_dir, "effb2.pt")
+        if not os.path.exists(model_path):
+            state_dict = load_state_dict_from_url(
+                'https://github.com/richermans/HEAR2021_EfficientLatent/releases/download/v0.0.1/effb2.pt',
+                progress=True,
+                model_dir=model_dir)
+        else:
+            state_dict = torch.load(model_path)
+        del_keys = [key for key in state_dict if key.startswith("front_end")]
+        for key in del_keys:
+            del state_dict[key]
+        model.eff_net.load_state_dict(state_dict)
+    return model
+
+
+class MBConvBlockPruned(MBConvBlock):
+
+    def __init__(self, block_args, global_params, image_size=None, prune_ratio=0.5, prune_se=True):
+        super(MBConvBlock, self).__init__()
+        self._block_args = block_args
+        self._bn_mom = 1 - global_params.batch_norm_momentum  # pytorch's difference from tensorflow
+        self._bn_eps = global_params.batch_norm_epsilon
+        self.has_se = (self._block_args.se_ratio is not None) and (0 < self._block_args.se_ratio <= 1)
+        self.id_skip = block_args.id_skip  # whether to use skip connection and drop connect
+
+        # Expansion phase (Inverted Bottleneck)
+        inp = self._block_args.input_filters  # number of input channels
+        oup = self._block_args.input_filters * self._block_args.expand_ratio  # number of output channels
+        if self._block_args.expand_ratio != 1:
+            oup = int(oup * (1 - prune_ratio))
+            Conv2d = get_same_padding_conv2d(image_size=image_size)
+            self._expand_conv = Conv2d(in_channels=inp, out_channels=oup, kernel_size=1, bias=False)
+            self._bn0 = nn.BatchNorm2d(num_features=oup, momentum=self._bn_mom, eps=self._bn_eps)
+            # image_size = calculate_output_image_size(image_size, 1) <-- this wouldn't modify image_size
+
+        # Depthwise convolution phase
+        k = self._block_args.kernel_size
+        s = self._block_args.stride
+        Conv2d = get_same_padding_conv2d(image_size=image_size)
+        self._depthwise_conv = Conv2d(
+            in_channels=oup, out_channels=oup, groups=oup,  # groups makes it depthwise
+            kernel_size=k, stride=s, bias=False)
+        self._bn1 = nn.BatchNorm2d(num_features=oup, momentum=self._bn_mom, eps=self._bn_eps)
+        image_size = calculate_output_image_size(image_size, s)
+
+        # Squeeze and Excitation layer, if desired
+        if self.has_se:
+            Conv2d = get_same_padding_conv2d(image_size=(1, 1))
+            num_squeezed_channels = max(1, int(self._block_args.input_filters * self._block_args.se_ratio))
+            if prune_se:
+                num_squeezed_channels = int(num_squeezed_channels * (1 - prune_ratio))
+            self._se_reduce = Conv2d(in_channels=oup, out_channels=num_squeezed_channels, kernel_size=1)
+            self._se_expand = Conv2d(in_channels=num_squeezed_channels, out_channels=oup, kernel_size=1)
+
+        # Pointwise convolution phase
+        final_oup = self._block_args.output_filters
+        Conv2d = get_same_padding_conv2d(image_size=image_size)
+        self._project_conv = Conv2d(in_channels=oup, out_channels=final_oup, kernel_size=1, bias=False)
+        self._bn2 = nn.BatchNorm2d(num_features=final_oup, momentum=self._bn_mom, eps=self._bn_eps)
+        self._swish = MemoryEfficientSwish()
+
+
+class EfficientNetB2Pruned(EfficientNet):
+
+    def __init__(self, blocks_args=None, global_params=None,
+                 prune_start_layer=0, prune_ratio=0.5, prune_se=True):
+        super(EfficientNet, self).__init__()
+        assert isinstance(blocks_args, list), 'blocks_args should be a list'
+        assert len(blocks_args) > 0, 'block args must be greater than 0'
+        self._global_params = global_params
+        self._blocks_args = blocks_args
+
+        # Batch norm parameters
+        bn_mom = 1 - self._global_params.batch_norm_momentum
+        bn_eps = self._global_params.batch_norm_epsilon
+
+        # Get stem static or dynamic convolution depending on image size
+        image_size = global_params.image_size
+        Conv2d = get_same_padding_conv2d(image_size=image_size)
+
+        n_build_blks = 0
+        # Stem
+        in_channels = 1  # spectrogram
+
+        p = 0.0 if n_build_blks < prune_start_layer else prune_ratio
+        out_channels = round_filters(32 * (1 - p),
+                                     self._global_params)  # number of output channels
+        self._conv_stem = Conv2d(in_channels, out_channels, kernel_size=3, stride=2, bias=False)
+        self._bn0 = nn.BatchNorm2d(num_features=out_channels, momentum=bn_mom, eps=bn_eps)
+        image_size = calculate_output_image_size(image_size, 2)
+        n_build_blks += 1
+
+        # Build blocks
+        self._blocks = nn.ModuleList([])
+        for block_args in self._blocks_args:
+
+            p = 0.0 if n_build_blks < prune_start_layer else prune_ratio
+            orig_input_filters = block_args.input_filters
+            # Update block input and output filters based on depth multiplier.
+            block_args = block_args._replace(
+                input_filters=round_filters(
+                    block_args.input_filters * (1 - p),
+                    self._global_params),
+                output_filters=round_filters(
+                    block_args.output_filters * (1 - p),
+                    self._global_params),
+                num_repeat=round_repeats(block_args.num_repeat, self._global_params)
+            )
+
+            if n_build_blks == prune_start_layer:
+                block_args = block_args._replace(input_filters=round_filters(
+                    orig_input_filters,
+                    self._global_params)
+                )
+
+            # The first block needs to take care of stride and filter size increase.
+            self._blocks.append(MBConvBlockPruned(block_args, self._global_params,
+                                                  image_size=image_size, prune_ratio=p,
+                                                  prune_se=prune_se))
+            n_build_blks += 1
+
+            image_size = calculate_output_image_size(image_size, block_args.stride)
+            if block_args.num_repeat > 1:  # modify block_args to keep same output size
+                block_args = block_args._replace(input_filters=block_args.output_filters, stride=1)
+            for _ in range(block_args.num_repeat - 1):
+                self._blocks.append(MBConvBlockPruned(block_args,
+                                                      self._global_params,
+                                                      image_size=image_size,
+                                                      prune_ratio=p,
+                                                      prune_se=prune_se))
+                # image_size = calculate_output_image_size(image_size, block_args.stride)  # stride = 1
+
+        # Head
+        in_channels = block_args.output_filters  # output of final block
+        p = 0.0 if n_build_blks < prune_start_layer else prune_ratio
+        out_channels = round_filters(1280 * (1 - p), self._global_params)
+        Conv2d = get_same_padding_conv2d(image_size=image_size)
+        self._conv_head = Conv2d(in_channels, out_channels, kernel_size=1, bias=False)
+        self._bn1 = nn.BatchNorm2d(num_features=out_channels, momentum=bn_mom, eps=bn_eps)
+
+        # Final linear layer
+        self._avg_pooling = nn.AdaptiveAvgPool2d(1)
+        if self._global_params.include_top:
+            self._dropout = nn.Dropout(self._global_params.dropout_rate)
+            self._fc = nn.Linear(out_channels, self._global_params.num_classes)
+
+        # set activation to memory efficient swish by default
+        self._swish = MemoryEfficientSwish()
+
+
+def get_pruned_model(pretrained: bool = True,
+                     prune_ratio: float = 0.5,
+                     prune_start_layer: int = 0,
+                     prune_se: bool = True,
+                     prune_method: str = "operator_norm") -> _EffiNet:
+    
+    import captioning.models.conv_filter_pruning as pruning_lib
+
+    blocks_args, global_params = efficientnet_utils.get_model_params(
+        'efficientnet-b2', {'include_top': False})
+    # print("num blocks: ", len(blocks_args))
+    # print("block args: ")
+    # for block_arg in blocks_args:
+    #     print(block_arg)
+    model = _EffiNet(blocks_args=blocks_args,
+                     global_params=global_params,
+                     prune_start_layer=prune_start_layer,
+                     prune_se=prune_se,
+                     prune_ratio=prune_ratio)
+    
+    if prune_method == "operator_norm":
+        filter_pruning = pruning_lib.operator_norm_pruning
+    elif prune_method == "interspeech":
+        filter_pruning = pruning_lib.cs_interspeech
+    elif prune_method == "iclr_l1":
+        filter_pruning = pruning_lib.iclr_l1
+    elif prune_method == "iclr_gm":
+        filter_pruning = pruning_lib.iclr_gm
+    elif prune_method == "cs_waspaa":
+        filter_pruning = pruning_lib.cs_waspaa
+        
+
+    if isinstance(pretrained, str):
+        ckpt = torch.load(pretrained, "cpu")
+        state_dict = {}
+        for key in ckpt["model"].keys():
+            if key.startswith("model.encoder.backbone"):
+                state_dict[key[len("model.encoder.backbone.eff_net."):]] = ckpt["model"][key]
+    elif isinstance(pretrained, bool):
+        model_path = os.path.join(model_dir, "effb2.pt")
+        if not os.path.exists(model_path):
+            state_dict = load_state_dict_from_url(
+                'https://github.com/richermans/HEAR2021_EfficientLatent/releases/download/v0.0.1/effb2.pt',
+                progress=True,
+                model_dir=model_dir)
+        else:
+            state_dict = torch.load(model_path)
+        del_keys = [key for key in state_dict if key.startswith("front_end")]
+        for key in del_keys:
+            del state_dict[key]
+    
+    # load pretrained model with corresponding filters
+    # rule:
+    # * depthwise_conv: in_ch_idx = out_ch_idx = prev_conv_idx
+    mod_dep_path = [
+        "_conv_stem",
+    ]
+    conv_to_bn = {"_conv_stem": "_bn0"}
+    for i in range(2):
+        mod_dep_path.extend([
+            f"_blocks.{i}._depthwise_conv",
+            f"_blocks.{i}._se_reduce",
+            f"_blocks.{i}._se_expand",
+            f"_blocks.{i}._project_conv", 
+        ])
+        conv_to_bn[f"_blocks.{i}._depthwise_conv"] = f"_blocks.{i}._bn1"
+        conv_to_bn[f"_blocks.{i}._project_conv"] = f"_blocks.{i}._bn2"
+    
+    for i in range(2, 23):
+        mod_dep_path.extend([
+            f"_blocks.{i}._expand_conv",
+            f"_blocks.{i}._depthwise_conv", 
+            f"_blocks.{i}._se_reduce",
+            f"_blocks.{i}._se_expand",
+            f"_blocks.{i}._project_conv"
+        ])
+        conv_to_bn[f"_blocks.{i}._expand_conv"] = f"_blocks.{i}._bn0"
+        conv_to_bn[f"_blocks.{i}._depthwise_conv"] = f"_blocks.{i}._bn1"
+        conv_to_bn[f"_blocks.{i}._project_conv"] = f"_blocks.{i}._bn2"
+
+    mod_dep_path.append("_conv_head") 
+    conv_to_bn["_conv_head"] = "_bn1"
+
+    # print(mod_dep_path)
+    # print(conv_to_bn)
+    
+    key_to_w_b_idx = {}
+    model_dict = model.eff_net.state_dict()
+    for conv_key in tqdm(mod_dep_path):
+        weight = state_dict[f"{conv_key}.weight"]
+        ptr_n_filter = weight.size(0)
+        model_n_filter = model_dict[f"{conv_key}.weight"].size(0)
+        if model_n_filter < ptr_n_filter:
+            key_to_w_b_idx[conv_key] = filter_pruning(weight.numpy())[:model_n_filter]
+        else:
+            key_to_w_b_idx[conv_key] = slice(None)
+
+    pruned_state_dict = {}
+    for conv_key, prev_conv_key in zip(mod_dep_path, [None] + mod_dep_path[:-1]):
+    
+        for sub_key in ["weight", "bias"]: # adjust the conv layer
+            cur_key = f"{conv_key}.{sub_key}"
+
+            if cur_key not in state_dict:
+                continue
+
+            if prev_conv_key is None or conv_key.endswith("_depthwise_conv"):
+                conv_in_idx = slice(None)
+            else:
+                conv_in_idx = key_to_w_b_idx[prev_conv_key]
+
+            # the first pruned layer
+            if model_dict[cur_key].ndim > 1 and model_dict[cur_key].size(1) == state_dict[cur_key].size(1):
+                conv_in_idx = slice(None)
+            
+            if conv_key.endswith("_depthwise_conv"):
+                conv_out_idx = key_to_w_b_idx[prev_conv_key]
+            else:
+                conv_out_idx = key_to_w_b_idx[conv_key]
+            
+            # if conv_key == "_blocks.16._se_reduce":
+            #     print(len(conv_out_idx), len(conv_in_idx))
+
+            if sub_key == "weight":
+                pruned_state_dict[cur_key] = state_dict[cur_key][
+                    conv_out_idx, ...][:, conv_in_idx, ...]
+            else:
+                pruned_state_dict[cur_key] = state_dict[cur_key][
+                    conv_out_idx, ...]
+
+        if conv_key in conv_to_bn: # adjust the corresponding bn layer
+            for sub_key in ["weight", "bias", "running_mean", "running_var"]:
+                cur_key = f"{conv_to_bn[conv_key]}.{sub_key}"
+                if cur_key not in state_dict:
+                    continue
+                pruned_state_dict[cur_key] = state_dict[cur_key][
+                    key_to_w_b_idx[conv_key], ...]
+    
+    model.eff_net.load_state_dict(pruned_state_dict)
+
+    return model

models/kd_wrapper.py

@@ -0,0 +1,226 @@
+from typing import Dict
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import repeat
+
+from models.base import CaptionMetaMixin
+from utils.model_util import init
+
+
+class WmlEncoderKdWrapper(nn.Module, CaptionMetaMixin):
+
+    def __init__(self,
+                 model: nn.Module,
+                 shared_dim: int,
+                 tchr_layer_to_dims: Dict[str, int],
+                 loss_type: str = "mse",):
+        super().__init__()
+        self.model = model
+        self.tchr_layers = list(tchr_layer_to_dims.keys())
+        self.stdnt_qv_proj = nn.Linear(model.encoder.fc_emb_size,
+                                       2 * shared_dim)
+        self.stdnt_qv_proj.apply(init)
+        for layer, dim in tchr_layer_to_dims.items():
+            self.add_module(f'tchr_kv_proj_{layer}', nn.Linear(dim, 2 * shared_dim))
+            getattr(self, f'tchr_kv_proj_{layer}').apply(init)
+        if loss_type == "mse":
+            self.loss_fn = nn.MSELoss(reduction="none")
+    
+    def forward(self, input_dict: Dict):
+        output_dict = self.model(input_dict)
+        if "tchr_output" in input_dict:
+            stdnt_emb = output_dict["fc_emb"]
+            stdnt_qv = self.stdnt_qv_proj(stdnt_emb)
+            stdnt_q, stdnt_v = torch.chunk(stdnt_qv, 2, dim=-1)
+            
+            tchr_output = input_dict["tchr_output"]
+            layer_ks, layer_vs = [], []
+            for layer in self.tchr_layers:
+                layer_kv = getattr(self, f'tchr_kv_proj_{layer}')(tchr_output[layer])
+                layer_k, layer_v = torch.chunk(layer_kv, 2, dim=-1)
+                layer_ks.append(layer_k)
+                layer_vs.append(layer_v)
+            layer_ks = torch.stack(layer_ks, dim=1)
+            layer_vs = torch.stack(layer_vs, dim=1)
+            weights = torch.softmax(stdnt_q.unsqueeze(1) @ layer_ks.transpose(1, 2), dim=-1)
+            stdnt_v = repeat(stdnt_v, 'b d -> b n d', n=len(self.tchr_layers))
+            loss = self.loss_fn(stdnt_v, layer_vs).mean(dim=-1, keepdim=True)
+            loss = (weights @ loss).mean()
+            output_dict["enc_kd_loss"] = loss
+        return output_dict
+
+
+class MseEncoderKdWrapper(nn.Module, CaptionMetaMixin):
+
+    def __init__(self,
+                 model: nn.Module,
+                 shared_dim: int,
+                 tchr_dim: int,
+                 use_tchr_proj: bool = True,
+                 l2_norm: bool = False,
+                 ):
+        super().__init__()
+        self.model = model
+        self.use_tchr_proj = use_tchr_proj
+        if not use_tchr_proj:
+            assert shared_dim == tchr_dim
+        self.tchr_dim = tchr_dim
+        self.l2_norm = l2_norm
+        if hasattr(model, "encoder"):
+            self.stdnt_proj = nn.Linear(model.encoder.fc_emb_size,
+                                        shared_dim)
+        else:
+            self.stdnt_proj = nn.Linear(model.fc_emb_size,
+                                        shared_dim)
+        self.stdnt_proj.apply(init)
+        if use_tchr_proj:
+            self.tchr_proj = nn.Linear(tchr_dim, shared_dim)
+            self.tchr_proj.apply(init)
+        else:
+            self.tchr_proj = nn.Identity()
+
+    def forward(self, input_dict: Dict):
+        unsup = input_dict.get("unsup", False)
+        if unsup is False:
+            if self.use_tchr_proj:
+                output_dict = self.model(input_dict)
+                stdnt_emb = output_dict["fc_emb"]
+            else:
+                encoder_output = self.model.encoder(input_dict)
+                stdnt_emb = encoder_output["fc_emb"]
+                encoder_output["fc_emb"] = self.stdnt_proj(encoder_output["fc_emb"])
+                encoder_output["attn_emb"] = self.stdnt_proj(encoder_output["attn_emb"])
+                output_dict = self.model.forward_decoder(input_dict, encoder_output)
+        else:
+            output_dict = self.model.encoder(input_dict)
+            stdnt_emb = output_dict["fc_emb"]
+        if "tchr_output" in input_dict:
+            stdnt_emb = self.stdnt_proj(stdnt_emb)
+            tchr_emb = input_dict["tchr_output"]["embedding"]
+            thcr_emb = self.tchr_proj(tchr_emb)
+
+            if self.l2_norm:
+                stdnt_emb = F.normalize(stdnt_emb, dim=-1)
+                thcr_emb = F.normalize(thcr_emb, dim=-1)
+
+            loss = F.mse_loss(stdnt_emb, thcr_emb)
+            output_dict["enc_kd_loss"] = loss
+        return output_dict
+
+
+class ContraEncoderKdWrapper(nn.Module, CaptionMetaMixin):
+
+    def __init__(self,
+                 model: nn.Module,
+                 shared_dim: int,
+                 tchr_dim: int,
+                 ):
+        super().__init__()
+        self.model = model
+        self.tchr_dim = tchr_dim
+        if hasattr(model, "encoder"):
+            self.stdnt_proj = nn.Linear(model.encoder.fc_emb_size,
+                                        shared_dim)
+        else:
+            self.stdnt_proj = nn.Linear(model.fc_emb_size,
+                                        shared_dim)
+        self.stdnt_proj.apply(init)
+        self.tchr_proj = nn.Linear(tchr_dim, shared_dim)
+        self.tchr_proj.apply(init)
+        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+
+    def forward(self, input_dict: Dict):
+        unsup = input_dict.get("unsup", False)
+        if unsup is False:
+            output_dict = self.model(input_dict)
+        else:
+            output_dict = self.model.encoder(input_dict)
+        if "tchr_output" in input_dict:
+            stdnt_emb = output_dict["fc_emb"]
+            stdnt_emb = self.stdnt_proj(stdnt_emb)
+            tchr_emb = input_dict["tchr_output"]["embedding"]
+            thcr_emb = self.tchr_proj(tchr_emb)
+
+            stdnt_emb = F.normalize(stdnt_emb, dim=-1)
+            thcr_emb = F.normalize(thcr_emb, dim=-1)
+
+            unscaled_logit = stdnt_emb @ thcr_emb.transpose(0, 1)
+            logit = self.logit_scale * unscaled_logit
+            label = torch.arange(logit.shape[0]).to(logit.device)
+            loss1 = F.cross_entropy(logit, label)
+            loss2 = F.cross_entropy(logit.transpose(0, 1), label)
+            loss = (loss1 + loss2) / 2
+            output_dict["enc_kd_loss"] = loss
+        return output_dict
+
+
+class ContraMseEncoderKdWrapper(nn.Module, CaptionMetaMixin):
+
+    def __init__(self,
+                 model: nn.Module,
+                 shared_dim: int,
+                 tchr_dim: int,
+                 use_tchr_proj: bool = True,
+                 l2_norm: bool = False,
+                 ):
+        super().__init__()
+        self.model = model
+        self.use_tchr_proj = use_tchr_proj
+        if not use_tchr_proj:
+            assert shared_dim == tchr_dim
+        self.tchr_dim = tchr_dim
+        self.l2_norm = l2_norm
+        if hasattr(model, "encoder"):
+            self.stdnt_proj = nn.Linear(model.encoder.fc_emb_size,
+                                        shared_dim)
+        else:
+            self.stdnt_proj = nn.Linear(model.fc_emb_size,
+                                        shared_dim)
+        self.stdnt_proj.apply(init)
+        if use_tchr_proj:
+            self.tchr_proj = nn.Linear(tchr_dim, shared_dim)
+            self.tchr_proj.apply(init)
+        else:
+            self.tchr_proj = nn.Identity()
+        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+
+    def forward(self, input_dict: Dict):
+        unsup = input_dict.get("unsup", False)
+        if unsup is False:
+            if self.use_tchr_proj:
+                output_dict = self.model(input_dict)
+                stdnt_emb = output_dict["fc_emb"]
+            else:
+                encoder_output = self.model.encoder(input_dict)
+                stdnt_emb = encoder_output["fc_emb"]
+                encoder_output["fc_emb"] = self.stdnt_proj(encoder_output["fc_emb"])
+                encoder_output["attn_emb"] = self.stdnt_proj(encoder_output["attn_emb"])
+                output_dict = self.model.forward_decoder(input_dict, encoder_output)
+        else:
+            output_dict = self.model.encoder(input_dict)
+            stdnt_emb = output_dict["fc_emb"]
+        if "tchr_output" in input_dict:
+            stdnt_emb = self.stdnt_proj(stdnt_emb)
+            tchr_emb = input_dict["tchr_output"]["embedding"]
+            thcr_emb = self.tchr_proj(tchr_emb)
+
+            if self.l2_norm:
+                stdnt_emb = F.normalize(stdnt_emb, dim=-1)
+                thcr_emb = F.normalize(thcr_emb, dim=-1)
+
+            mse_loss = F.mse_loss(stdnt_emb, thcr_emb)
+
+            stdnt_emb = F.normalize(stdnt_emb, dim=-1)
+            thcr_emb = F.normalize(thcr_emb, dim=-1)
+            unscaled_logit = stdnt_emb @ thcr_emb.transpose(0, 1)
+            logit = self.logit_scale * unscaled_logit
+            label = torch.arange(logit.shape[0]).to(logit.device)
+            loss1 = F.cross_entropy(logit, label)
+            loss2 = F.cross_entropy(logit.transpose(0, 1), label)
+            cntr_loss = (loss1 + loss2) / 2
+            output_dict["enc_kd_loss"] = mse_loss + cntr_loss
+
+        return output_dict

models/transformer_decoder.py

@@ -0,0 +1,214 @@
+import math
+
+import torch
+import torch.nn as nn
+
+from models import BaseDecoder
+from utils.model_util import generate_length_mask, PositionalEncoding
+from utils.train_util import merge_load_state_dict
+
+
+class TransformerDecoder(BaseDecoder):
+
+    def __init__(self,
+                 emb_dim,
+                 vocab_size,
+                 fc_emb_dim,
+                 attn_emb_dim,
+                 dropout,
+                 freeze=False,
+                 tie_weights=False,
+                 **kwargs):
+        super().__init__(emb_dim, vocab_size, fc_emb_dim, attn_emb_dim,
+                         dropout=dropout, tie_weights=tie_weights)
+        self.d_model = emb_dim
+        self.nhead = kwargs.get("nhead", self.d_model // 64)
+        self.nlayers = kwargs.get("nlayers", 2)
+        self.dim_feedforward = kwargs.get("dim_feedforward", self.d_model * 4)
+
+        self.pos_encoder = PositionalEncoding(self.d_model, dropout)
+        layer = nn.TransformerDecoderLayer(d_model=self.d_model,
+                                           nhead=self.nhead,
+                                           dim_feedforward=self.dim_feedforward,
+                                           dropout=dropout)
+        self.model = nn.TransformerDecoder(layer, self.nlayers)
+        self.classifier = nn.Linear(self.d_model, vocab_size, bias=False)
+        if tie_weights:
+            self.classifier.weight = self.word_embedding.weight
+        self.attn_proj = nn.Sequential(
+            nn.Linear(self.attn_emb_dim, self.d_model),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.LayerNorm(self.d_model)
+        )
+        self.init_params()
+
+        self.freeze = freeze
+        if freeze:
+            for p in self.parameters():
+                p.requires_grad = False
+
+    def init_params(self):
+        for p in self.parameters():
+            if p.dim() > 1:
+                nn.init.xavier_uniform_(p)
+    
+    def load_pretrained(self, pretrained, output_fn):
+        checkpoint = torch.load(pretrained, map_location="cpu")
+
+        if "model" in checkpoint:
+            checkpoint = checkpoint["model"]
+            if next(iter(checkpoint)).startswith("decoder."):
+                state_dict = {}
+                for k, v in checkpoint.items():
+                    state_dict[k[8:]] = v
+
+        loaded_keys = merge_load_state_dict(state_dict, self, output_fn)
+        if self.freeze:
+            for name, param in self.named_parameters():
+                if name in loaded_keys:
+                    param.requires_grad = False
+                else:
+                    param.requires_grad = True
+
+
+    def generate_square_subsequent_mask(self, max_length):
+        mask = (torch.triu(torch.ones(max_length, max_length)) == 1).transpose(0, 1)
+        mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
+        return mask
+
+    def forward(self, input_dict):
+        word = input_dict["word"]
+        attn_emb = input_dict["attn_emb"]
+        attn_emb_len = input_dict["attn_emb_len"]
+        cap_padding_mask = input_dict["cap_padding_mask"]
+
+        p_attn_emb = self.attn_proj(attn_emb)
+        p_attn_emb = p_attn_emb.transpose(0, 1) # [T_src, N, emb_dim]
+        word = word.to(attn_emb.device)
+        embed = self.in_dropout(self.word_embedding(word)) * math.sqrt(self.emb_dim) # [N, T, emb_dim]
+        embed = embed.transpose(0, 1) # [T, N, emb_dim]
+        embed = self.pos_encoder(embed)
+
+        tgt_mask = self.generate_square_subsequent_mask(embed.size(0)).to(attn_emb.device)
+        memory_key_padding_mask = ~generate_length_mask(attn_emb_len, attn_emb.size(1)).to(attn_emb.device)
+        output = self.model(embed, p_attn_emb, tgt_mask=tgt_mask,
+                            tgt_key_padding_mask=cap_padding_mask,
+                            memory_key_padding_mask=memory_key_padding_mask)
+        output = output.transpose(0, 1)
+        output = {
+            "embed": output,
+            "logit": self.classifier(output),
+        }
+        return output
+
+
+class M2TransformerDecoder(BaseDecoder):
+
+    def __init__(self, vocab_size, fc_emb_dim, attn_emb_dim, dropout=0.1, **kwargs):
+        super().__init__(attn_emb_dim, vocab_size, fc_emb_dim, attn_emb_dim, dropout=dropout,)
+        try:
+            from m2transformer.models.transformer import MeshedDecoder
+        except:
+            raise ImportError("meshed-memory-transformer not installed; please run `pip install git+https://github.com/ruotianluo/meshed-memory-transformer.git`")
+        del self.word_embedding
+        del self.in_dropout
+
+        self.d_model = attn_emb_dim
+        self.nhead = kwargs.get("nhead", self.d_model // 64)
+        self.nlayers = kwargs.get("nlayers", 2)
+        self.dim_feedforward = kwargs.get("dim_feedforward", self.d_model * 4)
+        self.model = MeshedDecoder(vocab_size, 100, self.nlayers, 0,
+                                   d_model=self.d_model,
+                                   h=self.nhead,
+                                   d_ff=self.dim_feedforward,
+                                   dropout=dropout)
+        self.init_params()
+
+    def init_params(self):
+        for p in self.parameters():
+            if p.dim() > 1:
+                nn.init.xavier_uniform_(p)
+
+    def forward(self, input_dict):
+        word = input_dict["word"]
+        attn_emb = input_dict["attn_emb"]
+        attn_emb_mask = input_dict["attn_emb_mask"]
+        word = word.to(attn_emb.device)
+        embed, logit = self.model(word, attn_emb, attn_emb_mask)
+        output = {
+            "embed": embed,
+            "logit": logit,
+        }
+        return output
+
+
+class EventTransformerDecoder(TransformerDecoder):
+
+    def forward(self, input_dict):
+        word = input_dict["word"] # index of word embeddings
+        attn_emb = input_dict["attn_emb"]
+        attn_emb_len = input_dict["attn_emb_len"]
+        cap_padding_mask = input_dict["cap_padding_mask"]
+        event_emb = input_dict["event"] # [N, emb_dim]
+
+        p_attn_emb = self.attn_proj(attn_emb)
+        p_attn_emb = p_attn_emb.transpose(0, 1) # [T_src, N, emb_dim]
+        word = word.to(attn_emb.device)
+        embed = self.in_dropout(self.word_embedding(word)) * math.sqrt(self.emb_dim) # [N, T, emb_dim]
+
+        embed = embed.transpose(0, 1) # [T, N, emb_dim]
+        embed += event_emb
+        embed = self.pos_encoder(embed)
+
+        tgt_mask = self.generate_square_subsequent_mask(embed.size(0)).to(attn_emb.device)
+        memory_key_padding_mask = ~generate_length_mask(attn_emb_len, attn_emb.size(1)).to(attn_emb.device)
+        output = self.model(embed, p_attn_emb, tgt_mask=tgt_mask,
+                            tgt_key_padding_mask=cap_padding_mask,
+                            memory_key_padding_mask=memory_key_padding_mask)
+        output = output.transpose(0, 1)
+        output = {
+            "embed": output,
+            "logit": self.classifier(output),
+        }
+        return output
+
+
+class KeywordProbTransformerDecoder(TransformerDecoder):
+
+    def __init__(self, emb_dim, vocab_size, fc_emb_dim, attn_emb_dim,
+                 dropout, keyword_classes_num, **kwargs):
+        super().__init__(emb_dim, vocab_size, fc_emb_dim, attn_emb_dim,
+                         dropout, **kwargs)
+        self.keyword_proj = nn.Linear(keyword_classes_num, self.d_model)
+        self.word_keyword_norm = nn.LayerNorm(self.d_model)
+
+    def forward(self, input_dict):
+        word = input_dict["word"] # index of word embeddings
+        attn_emb = input_dict["attn_emb"]
+        attn_emb_len = input_dict["attn_emb_len"]
+        cap_padding_mask = input_dict["cap_padding_mask"]
+        keyword = input_dict["keyword"] # [N, keyword_classes_num]
+
+        p_attn_emb = self.attn_proj(attn_emb)
+        p_attn_emb = p_attn_emb.transpose(0, 1) # [T_src, N, emb_dim]
+        word = word.to(attn_emb.device)
+        embed = self.in_dropout(self.word_embedding(word)) * math.sqrt(self.emb_dim) # [N, T, emb_dim]
+
+        embed = embed.transpose(0, 1) # [T, N, emb_dim]
+        embed += self.keyword_proj(keyword)
+        embed = self.word_keyword_norm(embed)
+
+        embed = self.pos_encoder(embed)
+
+        tgt_mask = self.generate_square_subsequent_mask(embed.size(0)).to(attn_emb.device)
+        memory_key_padding_mask = ~generate_length_mask(attn_emb_len, attn_emb.size(1)).to(attn_emb.device)
+        output = self.model(embed, p_attn_emb, tgt_mask=tgt_mask,
+                            tgt_key_padding_mask=cap_padding_mask,
+                            memory_key_padding_mask=memory_key_padding_mask)
+        output = output.transpose(0, 1)
+        output = {
+            "embed": output,
+            "logit": self.classifier(output),
+        }
+        return output

models/transformer_model.py

@@ -0,0 +1,264 @@
+# -*- coding: utf-8 -*-
+import random
+import torch
+import torch.nn as nn
+
+from models.base import CaptionModel
+from utils.model_util import repeat_tensor
+import models.transformer_decoder
+
+
+class TransformerModel(CaptionModel):
+
+    def __init__(self, encoder: nn.Module, decoder: nn.Module, **kwargs):
+        if not hasattr(self, "compatible_decoders"):
+            self.compatible_decoders = (
+                models.transformer_decoder.TransformerDecoder,
+            )
+        super().__init__(encoder, decoder, **kwargs)
+
+    def seq_forward(self, input_dict):
+        cap = input_dict["cap"]
+        cap_padding_mask = (cap == self.pad_idx).to(cap.device)
+        cap_padding_mask = cap_padding_mask[:, :-1]
+        output = self.decoder(
+            {
+                "word": cap[:, :-1],
+                "attn_emb": input_dict["attn_emb"],
+                "attn_emb_len": input_dict["attn_emb_len"],
+                "cap_padding_mask": cap_padding_mask
+            }
+        )
+        return output
+
+    def prepare_decoder_input(self, input_dict, output):
+        decoder_input = {
+            "attn_emb": input_dict["attn_emb"],
+            "attn_emb_len": input_dict["attn_emb_len"]
+        }
+        t = input_dict["t"]
+        
+        ###############
+        # determine input word
+        ################
+        if input_dict["mode"] == "train" and random.random() < input_dict["ss_ratio"]: # training, scheduled sampling
+            word = input_dict["cap"][:, :t+1]
+        else:
+            start_word = torch.tensor([self.start_idx,] * input_dict["attn_emb"].size(0)).unsqueeze(1).long()
+            if t == 0:
+                word = start_word
+            else:
+                word = torch.cat((start_word, output["seq"][:, :t]), dim=-1)
+        # word: [N, T]
+        decoder_input["word"] = word
+
+        cap_padding_mask = (word == self.pad_idx).to(input_dict["attn_emb"].device)
+        decoder_input["cap_padding_mask"] = cap_padding_mask
+        return decoder_input
+
+    def prepare_beamsearch_decoder_input(self, input_dict, output_i):
+        decoder_input = {}
+        t = input_dict["t"]
+        i = input_dict["sample_idx"]
+        beam_size = input_dict["beam_size"]
+        ###############
+        # prepare attn embeds
+        ################
+        if t == 0:
+            attn_emb = repeat_tensor(input_dict["attn_emb"][i], beam_size)
+            attn_emb_len = repeat_tensor(input_dict["attn_emb_len"][i], beam_size)
+            output_i["attn_emb"] = attn_emb
+            output_i["attn_emb_len"] = attn_emb_len
+        decoder_input["attn_emb"] = output_i["attn_emb"]
+        decoder_input["attn_emb_len"] = output_i["attn_emb_len"]
+        ###############
+        # determine input word
+        ################
+        start_word = torch.tensor([self.start_idx,] * beam_size).unsqueeze(1).long()
+        if t == 0:
+            word = start_word
+        else:
+            word = torch.cat((start_word, output_i["seq"]), dim=-1)
+        decoder_input["word"] = word
+        cap_padding_mask = (word == self.pad_idx).to(input_dict["attn_emb"].device)
+        decoder_input["cap_padding_mask"] = cap_padding_mask
+
+        return decoder_input
+
+
+class M2TransformerModel(CaptionModel):
+
+    def __init__(self, encoder: nn.Module, decoder: nn.Module, **kwargs):
+        if not hasattr(self, "compatible_decoders"):
+            self.compatible_decoders = (
+                models.transformer_decoder.M2TransformerDecoder,
+            )
+        super().__init__(encoder, decoder, **kwargs)
+        self.check_encoder_compatibility()
+
+    def check_encoder_compatibility(self):
+        assert isinstance(self.encoder, models.encoder.M2TransformerEncoder), \
+            f"only M2TransformerModel is compatible with {self.__class__.__name__}"
+
+    def seq_forward(self, input_dict):
+        cap = input_dict["cap"]
+        output = self.decoder(
+            {
+                "word": cap[:, :-1],
+                "attn_emb": input_dict["attn_emb"],
+                "attn_emb_mask": input_dict["attn_emb_mask"],
+            }
+        )
+        return output
+
+    def prepare_decoder_input(self, input_dict, output):
+        decoder_input = {
+            "attn_emb": input_dict["attn_emb"],
+            "attn_emb_mask": input_dict["attn_emb_mask"]
+        }
+        t = input_dict["t"]
+        
+        ###############
+        # determine input word
+        ################
+        if input_dict["mode"] == "train" and random.random() < input_dict["ss_ratio"]: # training, scheduled sampling
+            word = input_dict["cap"][:, :t+1]
+        else:
+            start_word = torch.tensor([self.start_idx,] * input_dict["attn_emb"].size(0)).unsqueeze(1).long()
+            if t == 0:
+                word = start_word
+            else:
+                word = torch.cat((start_word, output["seq"][:, :t]), dim=-1)
+        # word: [N, T]
+        decoder_input["word"] = word
+
+        return decoder_input
+
+    def prepare_beamsearch_decoder_input(self, input_dict, output_i):
+        decoder_input = {}
+        t = input_dict["t"]
+        i = input_dict["sample_idx"]
+        beam_size = input_dict["beam_size"]
+        ###############
+        # prepare attn embeds
+        ################
+        if t == 0:
+            attn_emb = repeat_tensor(input_dict["attn_emb"][i], beam_size)
+            attn_emb_mask = repeat_tensor(input_dict["attn_emb_mask"][i], beam_size)
+            output_i["attn_emb"] = attn_emb
+            output_i["attn_emb_mask"] = attn_emb_mask
+        decoder_input["attn_emb"] = output_i["attn_emb"]
+        decoder_input["attn_emb_mask"] = output_i["attn_emb_mask"]
+        ###############
+        # determine input word
+        ################
+        start_word = torch.tensor([self.start_idx,] * beam_size).unsqueeze(1).long()
+        if t == 0:
+            word = start_word
+        else:
+            word = torch.cat((start_word, output_i["seq"]), dim=-1)
+        decoder_input["word"] = word
+
+        return decoder_input
+
+
+class EventEncoder(nn.Module):
+    """
+    Encode the Label information in AudioCaps and AudioSet
+    """
+    def __init__(self, emb_dim, vocab_size=527):
+        super(EventEncoder, self).__init__()
+        self.label_embedding = nn.Parameter(
+            torch.randn((vocab_size, emb_dim)), requires_grad=True)
+        
+    def forward(self, word_idxs):
+        indices = word_idxs / word_idxs.sum(dim=1, keepdim=True)
+        embeddings = indices @ self.label_embedding
+        return embeddings
+
+
+class EventCondTransformerModel(TransformerModel):
+
+    def __init__(self, encoder: nn.Module, decoder: nn.Module, **kwargs):
+        if not hasattr(self, "compatible_decoders"):
+            self.compatible_decoders = (
+                models.transformer_decoder.EventTransformerDecoder,
+            )
+        super().__init__(encoder, decoder, **kwargs)
+        self.label_encoder = EventEncoder(decoder.emb_dim, 527)
+        self.train_forward_keys += ["events"]
+        self.inference_forward_keys += ["events"]
+
+    # def seq_forward(self, input_dict):
+        # cap = input_dict["cap"]
+        # cap_padding_mask = (cap == self.pad_idx).to(cap.device)
+        # cap_padding_mask = cap_padding_mask[:, :-1]
+        # output = self.decoder(
+            # {
+                # "word": cap[:, :-1],
+                # "attn_emb": input_dict["attn_emb"],
+                # "attn_emb_len": input_dict["attn_emb_len"],
+                # "cap_padding_mask": cap_padding_mask
+            # }
+        # )
+        # return output
+
+    def prepare_decoder_input(self, input_dict, output):
+        decoder_input = super().prepare_decoder_input(input_dict, output)
+        decoder_input["events"] = self.label_encoder(input_dict["events"])
+        return decoder_input
+
+    def prepare_beamsearch_decoder_input(self, input_dict, output_i):
+        decoder_input = super().prepare_beamsearch_decoder_input(input_dict, output_i)
+        t = input_dict["t"]
+        i = input_dict["sample_idx"]
+        beam_size = input_dict["beam_size"]
+        if t == 0:
+            output_i["events"] = repeat_tensor(self.label_encoder(input_dict["events"])[i], beam_size)
+        decoder_input["events"] = output_i["events"]
+        return decoder_input
+
+
+class KeywordCondTransformerModel(TransformerModel):
+
+    def __init__(self, encoder: nn.Module, decoder: nn.Module, **kwargs):
+        if not hasattr(self, "compatible_decoders"):
+            self.compatible_decoders = (
+                models.transformer_decoder.KeywordProbTransformerDecoder,
+            )
+        super().__init__(encoder, decoder, **kwargs)
+        self.train_forward_keys += ["keyword"]
+        self.inference_forward_keys += ["keyword"]
+
+    def seq_forward(self, input_dict):
+        cap = input_dict["cap"]
+        cap_padding_mask = (cap == self.pad_idx).to(cap.device)
+        cap_padding_mask = cap_padding_mask[:, :-1]
+        keyword = input_dict["keyword"]
+        output = self.decoder(
+            {
+                "word": cap[:, :-1],
+                "attn_emb": input_dict["attn_emb"],
+                "attn_emb_len": input_dict["attn_emb_len"],
+                "keyword": keyword,
+                "cap_padding_mask": cap_padding_mask
+            }
+        )
+        return output
+
+    def prepare_decoder_input(self, input_dict, output):
+        decoder_input = super().prepare_decoder_input(input_dict, output)
+        decoder_input["keyword"] = input_dict["keyword"]
+        return decoder_input
+
+    def prepare_beamsearch_decoder_input(self, input_dict, output_i):
+        decoder_input = super().prepare_beamsearch_decoder_input(input_dict, output_i)
+        t = input_dict["t"]
+        i = input_dict["sample_idx"]
+        beam_size = input_dict["beam_size"]
+        if t == 0:
+            output_i["keyword"] = repeat_tensor(input_dict["keyword"][i],
+                                                 beam_size)
+        decoder_input["keyword"] = output_i["keyword"]
+        return decoder_input
+

requirements.txt

@@ -0,0 +1,2 @@
+efficientnet_pytorch
+PyYAML

text_tokenizer.py

@@ -0,0 +1,107 @@
+import pickle
+from pathlib import Path
+
+import numpy as np
+from utils.train_util import pad_sequence
+
+
+class DictTokenizer:
+
+    def __init__(self,
+                 tokenizer_path: str = None,
+                 max_length: int = 20) -> None:
+        self.word2idx = {}
+        self.idx2word = {}
+        self.idx = 0
+        self.add_word("<pad>")
+        self.add_word("<start>")
+        self.add_word("<end>")
+        self.add_word("<unk>")
+        if tokenizer_path is not None and Path(tokenizer_path).exists():
+            state_dict = pickle.load(open(tokenizer_path, "rb"))
+            self.load_state_dict(state_dict)
+            self.loaded = True
+        else:
+            self.loaded = False
+        self.bos, self.eos = self.word2idx["<start>"], self.word2idx["<end>"]
+        self.pad = self.word2idx["<pad>"]
+        self.max_length = max_length
+
+    def add_word(self, word):
+        if not word in self.word2idx:
+            self.word2idx[word] = self.idx
+            self.idx2word[self.idx] = word
+            self.idx += 1
+
+    def encode_word(self, word):
+        if word in self.word2idx:
+            return self.word2idx[word]
+        else:
+            return self.word2idx["<unk>"]
+
+    def __call__(self, texts):
+        assert isinstance(texts, list), "the input must be List[str]"
+        batch_tokens = []
+        for text in texts:
+            tokens = [self.encode_word(token) for token in text.split()][:self.max_length]
+            tokens = [self.bos] + tokens + [self.eos]
+            tokens = np.array(tokens)
+            batch_tokens.append(tokens)
+        caps, cap_lens = pad_sequence(batch_tokens, self.pad)
+        return {
+            "cap": caps,
+            "cap_len": cap_lens
+        }
+
+    def decode(self, batch_token_ids):
+        output = []
+        for token_ids in batch_token_ids:
+            tokens = []
+            for token_id in token_ids:
+                if token_id == self.eos:
+                    break
+                elif token_id == self.bos:
+                    continue
+                tokens.append(self.idx2word[token_id])
+            output.append(" ".join(tokens))
+        return output
+
+    def __len__(self):
+        return len(self.word2idx)
+
+    def state_dict(self):
+        return self.word2idx
+    
+    def load_state_dict(self, state_dict):
+        self.word2idx = state_dict
+        self.idx2word = {idx: word for word, idx in self.word2idx.items()}
+        self.idx = len(self.word2idx)
+
+
+class HuggingfaceTokenizer:
+
+    def __init__(self,
+                 model_name_or_path,
+                 max_length) -> None:
+        from transformers import AutoTokenizer
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
+        self.max_length = max_length
+        self.bos, self.eos = self.tokenizer.bos_token_id, self.tokenizer.eos_token_id
+        self.pad = self.tokenizer.pad_token_id
+        self.loaded = True
+
+    def __call__(self, texts):
+        assert isinstance(texts, list), "the input must be List[str]"
+        batch_token_dict = self.tokenizer(texts,
+                                          padding=True,
+                                          truncation=True,
+                                          max_length=self.max_length,
+                                          return_tensors="pt")
+        batch_token_dict["cap"] = batch_token_dict["input_ids"]
+        cap_lens = batch_token_dict["attention_mask"].sum(dim=1)
+        cap_lens = cap_lens.numpy().astype(np.int32)
+        batch_token_dict["cap_len"] = cap_lens
+        return batch_token_dict
+
+    def decode(self, batch_token_ids):
+        return self.tokenizer.batch_decode(batch_token_ids, skip_special_tokens=True)

utils/model_util.py

@@ -0,0 +1,186 @@
+import math
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+from torch.nn.utils.rnn import PackedSequence, pack_padded_sequence, pad_packed_sequence
+
+
+def sort_pack_padded_sequence(input, lengths):
+    sorted_lengths, indices = torch.sort(lengths, descending=True)
+    tmp = pack_padded_sequence(input[indices], sorted_lengths.cpu(), batch_first=True)
+    inv_ix = indices.clone()
+    inv_ix[indices] = torch.arange(0,len(indices)).type_as(inv_ix)
+    return tmp, inv_ix
+
+def pad_unsort_packed_sequence(input, inv_ix):
+    tmp, _ = pad_packed_sequence(input, batch_first=True)
+    tmp = tmp[inv_ix]
+    return tmp
+
+def pack_wrapper(module, attn_feats, attn_feat_lens):
+    packed, inv_ix = sort_pack_padded_sequence(attn_feats, attn_feat_lens)
+    if isinstance(module, torch.nn.RNNBase):
+        return pad_unsort_packed_sequence(module(packed)[0], inv_ix)
+    else:
+        return pad_unsort_packed_sequence(PackedSequence(module(packed[0]), packed[1]), inv_ix)
+
+def generate_length_mask(lens, max_length=None):
+    lens = torch.as_tensor(lens)
+    N = lens.size(0)
+    if max_length is None:
+        max_length = max(lens)
+        if isinstance(max_length, torch.Tensor):
+            max_length = max_length.item()
+    idxs = torch.arange(max_length).repeat(N).view(N, max_length)
+    idxs = idxs.to(lens.device)
+    mask = (idxs < lens.view(-1, 1))
+    return mask
+
+def mean_with_lens(features, lens):
+    """
+    features: [N, T, ...] (assume the second dimension represents length)
+    lens: [N,]
+    """
+    lens = torch.as_tensor(lens)
+    if max(lens) != features.size(1):
+        max_length = features.size(1)
+        mask = generate_length_mask(lens, max_length)
+    else:
+        mask = generate_length_mask(lens)
+    mask = mask.to(features.device) # [N, T]
+
+    while mask.ndim < features.ndim:
+        mask = mask.unsqueeze(-1)
+    feature_mean = features * mask
+    feature_mean = feature_mean.sum(1)
+    while lens.ndim < feature_mean.ndim:
+        lens = lens.unsqueeze(1)
+    feature_mean = feature_mean / lens.to(features.device)
+    # feature_mean = features * mask.unsqueeze(-1)
+    # feature_mean = feature_mean.sum(1) / lens.unsqueeze(1).to(features.device)
+    return feature_mean
+
+def max_with_lens(features, lens):
+    """
+    features: [N, T, ...] (assume the second dimension represents length)
+    lens: [N,]
+    """
+    lens = torch.as_tensor(lens)
+    if max(lens) != features.size(1):
+        max_length = features.size(1)
+        mask = generate_length_mask(lens, max_length)
+    else:
+        mask = generate_length_mask(lens)
+    mask = mask.to(features.device) # [N, T]
+
+    feature_max = features.clone()
+    feature_max[~mask] = float("-inf")
+    feature_max, _ = feature_max.max(1)
+    return feature_max
+
+def repeat_tensor(x, n):
+    return x.unsqueeze(0).repeat(n, *([1] * len(x.shape)))
+
+def init(m, method="kaiming"):
+    if isinstance(m, (nn.Conv2d, nn.Conv1d)):
+        if method == "kaiming":
+            nn.init.kaiming_uniform_(m.weight)
+        elif method == "xavier":
+            nn.init.xavier_uniform_(m.weight)
+        else:
+            raise Exception(f"initialization method {method} not supported")
+        if m.bias is not None:
+            nn.init.constant_(m.bias, 0)
+    elif isinstance(m, (nn.BatchNorm2d, nn.BatchNorm1d)):
+        nn.init.constant_(m.weight, 1)
+        if m.bias is not None:
+            nn.init.constant_(m.bias, 0)
+    elif isinstance(m, nn.Linear):
+        if method == "kaiming":
+            nn.init.kaiming_uniform_(m.weight)
+        elif method == "xavier":
+            nn.init.xavier_uniform_(m.weight)
+        else:
+            raise Exception(f"initialization method {method} not supported")
+        if m.bias is not None:
+            nn.init.constant_(m.bias, 0)
+    elif isinstance(m, nn.Embedding):
+        if method == "kaiming":
+            nn.init.kaiming_uniform_(m.weight)
+        elif method == "xavier":
+            nn.init.xavier_uniform_(m.weight)
+        else:
+            raise Exception(f"initialization method {method} not supported")
+
+def compute_batch_score(decode_res,
+                        key2refs,
+                        keys,
+                        start_idx,
+                        end_idx,
+                        vocabulary,
+                        scorer):
+    """
+    Args:
+        decode_res: decoding results of model, [N, max_length]
+        key2refs: references of all samples, dict(<key> -> [ref_1, ref_2, ..., ref_n]
+        keys: keys of this batch, used to match decode results and refs
+    Return:
+        scores of this batch, [N,]
+    """
+
+    if scorer is None:
+        from pycocoevalcap.cider.cider import Cider
+        scorer = Cider()
+
+    hypothesis = {}
+    references = {}
+
+    for i in range(len(keys)):
+
+        if keys[i] in hypothesis.keys():
+            continue
+
+        # prepare candidate sentence
+        candidate = []
+        for w_t in decode_res[i]:
+            if w_t == start_idx:
+                continue
+            elif w_t == end_idx:
+                break
+            candidate.append(vocabulary.idx2word[w_t])
+
+        hypothesis[keys[i]] = [" ".join(candidate), ]
+
+        # prepare reference sentences
+        references[keys[i]] = key2refs[keys[i]]
+
+    score, scores = scorer.compute_score(references, hypothesis)
+    key2score = {key: scores[i] for i, key in enumerate(references.keys())}
+    results = np.zeros(decode_res.shape[0])
+    for i in range(decode_res.shape[0]):
+        results[i] = key2score[keys[i]]
+    return results 
+
+
+class PositionalEncoding(nn.Module):
+
+    def __init__(self, d_model, dropout=0.1, max_len=100):
+        super(PositionalEncoding, self).__init__()
+        self.dropout = nn.Dropout(p=dropout)
+
+        pe = torch.zeros(max_len, d_model)
+        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2).float() * \
+            (-math.log(10000.0) / d_model))
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        pe = pe.unsqueeze(0).transpose(0, 1)
+        # self.register_buffer("pe", pe)
+        self.register_parameter("pe", nn.Parameter(pe, requires_grad=False))
+
+    def forward(self, x):
+        # x: [T, N, E]
+        x = x + self.pe[:x.size(0), :]
+        return self.dropout(x)

utils/train_util.py

@@ -0,0 +1,117 @@
+import importlib
+import os
+import sys
+from typing import Callable, Dict, Union
+
+import numpy as np
+import yaml
+import torch
+
+
+def merge_a_into_b(a, b):
+    # merge dict a into dict b. values in a will overwrite b.
+    for k, v in a.items():
+        if isinstance(v, dict) and k in b:
+            assert isinstance(
+                b[k], dict
+            ), "Cannot inherit key '{}' from base!".format(k)
+            merge_a_into_b(v, b[k])
+        else:
+            b[k] = v
+
+
+def load_config(config_file):
+    with open(config_file, "r") as reader:
+        config = yaml.load(reader, Loader=yaml.FullLoader)
+    if "inherit_from" in config:
+        base_config_file = config["inherit_from"]
+        base_config_file = os.path.join(
+            os.path.dirname(config_file), base_config_file
+        )
+        assert not os.path.samefile(config_file, base_config_file), \
+            "inherit from itself"
+        base_config = load_config(base_config_file)
+        del config["inherit_from"]
+        merge_a_into_b(config, base_config)
+        return base_config
+    return config
+
+def get_cls_from_str(string, reload=False):
+    module_name, cls_name = string.rsplit(".", 1)
+    if reload:
+        module_imp = importlib.import_module(module_name)
+        importlib.reload(module_imp)
+    return getattr(importlib.import_module(module_name, package=None), cls_name)
+
+def init_obj_from_dict(config, **kwargs):
+    obj_args = config["args"].copy()
+    obj_args.update(kwargs)
+    for k in config:
+        if k not in ["type", "args"] and isinstance(config[k], dict) and k not in kwargs:
+            obj_args[k] = init_obj_from_dict(config[k])
+    try:
+        obj = get_cls_from_str(config["type"])(**obj_args)
+        return obj
+    except Exception as e:
+        print(f"Initializing {config} failed, detailed error stack: ")
+        raise e
+
+def init_model_from_config(config, print_fn=sys.stdout.write):
+    kwargs = {}
+    for k in config:
+        if k not in ["type", "args", "pretrained"]:
+            sub_model = init_model_from_config(config[k], print_fn)
+            if "pretrained" in config[k]:
+                load_pretrained_model(sub_model,
+                                      config[k]["pretrained"],
+                                      print_fn)
+            kwargs[k] = sub_model
+    model = init_obj_from_dict(config, **kwargs)
+    return model
+
+def merge_load_state_dict(state_dict,
+                          model: torch.nn.Module,
+                          output_fn: Callable = sys.stdout.write):
+    model_dict = model.state_dict()
+    pretrained_dict = {}
+    mismatch_keys = []
+    for key, value in state_dict.items():
+        if key in model_dict and model_dict[key].shape == value.shape:
+            pretrained_dict[key] = value
+        else:
+            mismatch_keys.append(key)
+    output_fn(f"Loading pre-trained model, with mismatched keys {mismatch_keys}")
+    model_dict.update(pretrained_dict)
+    model.load_state_dict(model_dict, strict=True)
+    return pretrained_dict.keys()
+
+
+def load_pretrained_model(model: torch.nn.Module,
+                          pretrained: Union[str, Dict],
+                          output_fn: Callable = sys.stdout.write):
+    if not isinstance(pretrained, dict) and not os.path.exists(pretrained):
+        output_fn(f"pretrained {pretrained} not exist!")
+        return
+    
+    if hasattr(model, "load_pretrained"):
+        model.load_pretrained(pretrained, output_fn)
+        return
+
+    if isinstance(pretrained, dict):
+        state_dict = pretrained
+    else:
+        state_dict = torch.load(pretrained, map_location="cpu")
+
+    if "model" in state_dict:
+        state_dict = state_dict["model"]
+    
+    merge_load_state_dict(state_dict, model, output_fn)
+
+def pad_sequence(data, pad_value=0):
+    if isinstance(data[0], (np.ndarray, torch.Tensor)):
+        data = [torch.as_tensor(arr) for arr in data]
+    padded_seq = torch.nn.utils.rnn.pad_sequence(data,
+                                                 batch_first=True,
+                                                 padding_value=pad_value)
+    length = np.array([x.shape[0] for x in data])
+    return padded_seq, length