Change to Hugging Face calling

app.py

@@ -1,25 +1,28 @@
-from pathlib import Path
-import argparse
 from functools import partial
 import gradio as gr
 import torch
 from torchaudio.functional import resample
+from transformers import AutoModel, PreTrainedTokenizerFast
 
-import utils.train_util as train_util
 
-
-def load_model(cfg,
-               ckpt_path,
+def load_model(model_name,
                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)
+    if model_name == "AudioCaps":
+        model = AutoModel.from_pretrained(
+            "wsntxxn/effb2-trm-audiocaps-captioning",
+            trust_remote_code=True
+        ).to(device)
+        tokenizer = PreTrainedTokenizerFast.from_pretrained(
+            "wsntxxn/audiocaps-simple-tokenizer"
+        )
+    elif model_name == "Clotho":
+        model = AutoModel.from_pretrained(
+            "wsntxxn/effb2-trm-clotho-captioning",
+            trust_remote_code=True
+        ).to(device)
+        tokenizer = PreTrainedTokenizerFast.from_pretrained(
+            "wsntxxn/clotho-simple-tokenizer"
+        )
     return model, tokenizer
 
 
@@ -34,19 +37,13 @@ def infer(file, runner):
         wav = wav.mean(1)
     wav = resample(wav, sr, runner.target_sr)
     wav_len = len(wav)
-    wav = wav.float().unsqueeze(0).to(runner.device)
-    input_dict = {
-        "mode": "inference",
-        "wav": wav,
-        "wav_len": [wav_len],
-        "specaug": False,
-        "sample_method": "beam",
-        "beam_size": 3,
-    }
+    wav = wav.float().unsqueeze(0)
     with torch.no_grad():
-        output_dict = runner.model(input_dict)
-        seq = output_dict["seq"].cpu().numpy()
-        cap = runner.tokenizer.decode(seq)[0]
+        word_idx = runner.model(
+            audio=wav,
+            audio_length=[wav_len]
+        )[0]
+        cap = runner.tokenizer.decode(word_idx, skip_special_tokens=True)
     return cap
 
 # def input_toggle(input_type):
@@ -59,16 +56,12 @@ class InferRunner:
 
     def __init__(self, model_name):
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        exp_dir = Path(f"./checkpoints/{model_name.lower()}")
-        cfg = train_util.load_config(exp_dir / "config.yaml")
-        self.model, self.tokenizer = load_model(cfg, exp_dir / "ckpt.pth", self.device)
-        self.target_sr = cfg["target_sr"]
+        self.model, self.tokenizer = load_model(model_name, self.device)
+        self.target_sr = self.model.config.sample_rate
     
     def change_model(self, model_name):
-        exp_dir = Path(f"./checkpoints/{model_name.lower()}")
-        cfg = train_util.load_config(exp_dir / "config.yaml")
-        self.model, self.tokenizer = load_model(cfg, exp_dir / "ckpt.pth", self.device)
-        self.target_sr = cfg["target_sr"]
+        self.model, self.tokenizer = load_model(model_name, self.device)
+        self.target_sr = self.model.config.sample_rate
 
 
 def change_model(radio):

checkpoints/audiocaps/ckpt.pth

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

checkpoints/audiocaps/config.yaml

@@ -1,30 +0,0 @@
-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

checkpoints/clotho/ckpt.pth

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

checkpoints/clotho/config.yaml

@@ -1,30 +0,0 @@
-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: 4368
-            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

@@ -1,92 +0,0 @@
-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

@@ -1,504 +0,0 @@
-# -*- 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

@@ -1,808 +0,0 @@
-# -*- coding: utf-8 -*-
-
-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

@@ -1,347 +0,0 @@
-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 = os.getcwd()
-
-
-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

@@ -1,226 +0,0 @@
-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

@@ -1,214 +0,0 @@
-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

@@ -1,264 +0,0 @@
-# -*- 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

@@ -1,4 +1,4 @@
+transformers
 efficientnet_pytorch
-PyYAML
 torchaudio
-einops
+einops

text_tokenizer.py

@@ -1,107 +0,0 @@
-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

@@ -1,186 +0,0 @@
-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

@@ -1,117 +0,0 @@
-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}\n")
-    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