Converge problem of a Region Caption Model

username: Anakin
cooked: <p>Hello,</p>
<p>I’m working on a region captioning task. I’m not sure what mistakes I might have done, as I see that the loss does not decrease forever. Could somebody give me a short review of my model. Thanks in advance.</p>
<pre><code class="lang-auto">import math

import torch.nn as nn
import torch.nn.functional as F
import torch
from torch.autograd import Variable
from .config import cfg
from os.path import join as opj
import json

# Language Criterion
class TemporalCrossEntropyCriterion(nn.Module):
    def __init__(self):
        super(TemporalCrossEntropyCriterion, self).__init__()
        self.lsm = nn.LogSoftmax(dim=1)
        
        # Whether to average over space and batch
        self.batch_average = True
        self.time_average = False
    
    def forward(self, input, target):
        # print("input.shape", input.shape)
        N, T, C = input.shape
        assert (target.dim() == 2 and target.size(0) == N and target.size(1) == T)
        null_mask = torch.eq(target, 0)
        target[null_mask] = 1

        ############time step case############
        # null_mask = torch.eq(target, 0)
        # target = target[:, :input.size(1)]
        # null_mask = null_mask[:, :input.size(1)]
        ############time step case############
        
        logprobs = self.lsm.forward(input)
        #losses = torch.gather(logprobs, 2, target.unsqueeze(-1)).mul(-1).squeeze(-1)
        losses = torch.gather(logprobs, 2, target.clone().view(N, T, -1)).mul(-1).squeeze(-1)
        losses[null_mask] = 0
        if self.batch_average:
            losses = torch.sum(losses) / N  # torch.sum(null_mask)
        if self.time_average:
            losses = torch.sum(losses) / T
        target[null_mask] = 0
        return losses

# Caption Model
class LSTM(nn.Module):
    def __init__(self, input_dim, hidden_dim, device="cuda"):
        super(LSTM, self).__init__()
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        D, H = input_dim, hidden_dim
        
        self.weight = nn.Parameter(torch.randn(D + H, 4 * H))
        self.bias = nn.Parameter(torch.randn(4 * H))
        
        self.hidden = torch.empty(0).to(device)
        self.cell = torch.empty(0).to(device)
        self.gates = torch.empty(0).to(device)
        
        self.h0 = torch.empty(0).to(device)
        self.c0 = torch.empty(0).to(device)
        
        self.remember_states = False
        self.device = device
    
    def reset(self, std):
        if not std:
            std = 1.0 / math.sqrt(self.hidden_dim + self.input_dim)
        
        self.bias.zero_()
        self.bias[self.hidden_dim, 2 * self.hidden_dim].fill_(1)
        self.weight.normal_(0, std)
    
    def reset_states(self):
        self.h0 = torch.empty(0).to(self.device)
        self.c0 = torch.empty(0).to(self.device)
    
    def check_dims(self, x, dims):
        assert (x.dim() == len(dims))
        for i, d in enumerate(dims):
            assert x.size(i) == d
    
    def unpack_input(self, input):
        c0, h0, x = None, None, None
        if type(input) == tuple and len(tuple) == 3:
            c0, h0, x = input
        elif type(input) == tuple and len(tuple) == 2:
            h0, x = input
        elif torch.is_tensor(input):
            x = input
        else:
            assert False, "invalid input"
        
        return c0, h0, x
    
    def get_sizes(self, input):
        c0, h0, x = self.unpack_input(input)
        N, T = x.size(0), x.size(1)
        H, D = self.hidden_dim, self.input_dim
        self.check_dims(x, [N, T, D])
        if h0:
            self.check_dims(h0, [N, H])
        if c0:
            self.check_dims(c0, [N, H])
        
        return N, T, D, H
    
    def lstm_step_forward(self, x, prev_h, prev_c, Wx, Wh, b):
        next_h, next_c = None, None

        H = prev_h.shape[1]
        
        gates = torch.addmm(b, x, Wx)

        gates.addmm(prev_h, Wh)
        
        # slice gate vector
        gate_i = gates[:, 0:H]
        gate_f = gates[:, H:2 * H]
        gate_o = gates[:, 2 * H:3 * H]
        gate_g = gates[:, 3 * H:]
        
        # activation functions applied to our 4 gates.
        input_gate = torch.sigmoid(gate_i)
        forget_gate = torch.sigmoid(gate_f)
        output_gate = torch.sigmoid(gate_o)
        block_input = torch.tanh(gate_g)
        
        # calculate next cell state
        next_c = (forget_gate * prev_c) + (input_gate * block_input)
        
        # calculate next hidden state
        next_h = output_gate * torch.tanh(next_c)
        
        return next_h, next_c
    
    def forward(self, input):
        c0, h0, x = self.unpack_input(input)
        N, T, D, H = self.get_sizes(input)
        
        if not c0:
            c0 = self.c0
            if c0.numel() == 0 or not self.remember_states:
                # c0.resize_(N, H).zero_().to(self.device)
                self.c0 = Variable(torch.zeros(N, H)).to(self.device)
                c0 = self.c0
            elif self.remember_states:
                # c0 = self.c0
                prev_N, prev_T = self.cell.shape
                assert prev_N == N, 'batch sizes must be constant to remember states'
                c0 = self.cell[:, :prev_T]
        
        if not h0:
            h0 = self.h0
            if h0.numel() == 0 or not self.remember_states:
                self.h0 = Variable(torch.zeros(N, H)).to(self.device)
                h0 = self.h0
            elif self.remember_states:
                prev_N, prev_T = self.hidden.shape
                assert prev_N == N, 'batch sizes must be the same to remember states'
                h0 = self.hidden[:, :prev_T]
        
        bias_expand = self.bias.view(1, 4 * H).expand(N, 4 * H).to(self.device)
        
        Wx = self.weight[:D, :]
        
        Wh = self.weight[D:D + H, :]
        
        self.hidden = Variable(torch.zeros(N, T, H).to(self.device))
        self.cell = Variable(torch.zeros(N, T, H).to(self.device))
        # h, c = self.output, self.cell
        
        prev_h, prev_c = h0, c0
        
        for t in range(T):
            prev_h, prev_c = self.lstm_step_forward(x[:, t, :], prev_h, prev_c, Wx, Wh, bias_expand)
            self.hidden[:, t, :] = prev_h
            self.cell[:, t, :] = prev_c
        
        if not self.training:
            self.hidden = self.hidden.squeeze(1)
            self.cell = self.cell.squeeze(1)
        
        return self.hidden


class ParallelTable(nn.Module):
    def __init__(self, image_input_dim=4096, image_embed_dim=512, rnn_size=512, vocab_size=512):
        super(ParallelTable, self).__init__()
        # For mapping from image vectors to word vectors
        self.image_encoder = nn.Sequential(
            nn.Linear(image_input_dim, image_embed_dim),
            nn.ReLU(True)
        )
        # For mapping word indices to word vectors
        self.lookup_table = nn.Embedding(vocab_size + 2, rnn_size)
    
    def forward(self, inputs):
        assert isinstance(inputs, tuple), "invalid format of inputs"
        feats, seqs = inputs
        encode_feats = self.image_encoder(feats)
        emd_words = self.lookup_table(seqs)
        return torch.cat([encode_feats.unsqueeze(1), emd_words], dim=1)


class RNN_LanguageModel(nn.Module):
    def __init__(self, opt=None, cfg=None):
        super(RNN_LanguageModel, self).__init__()
        
        vocab_pth = opj(cfg.DATASET.DATA_DIR, cfg.LANGUAGE_MODEL.VOCAB_FILE)
        with open(vocab_pth, "r") as fd:
            region_dicts = json.load(fd)
            cfg.LANGUAGE_MODEL.VOCAB_SIZE = len(region_dicts["vocab"])
            token_to_idx = region_dicts["token_to_idx"]
            self.idx_to_token = region_dicts["idx_to_token"]
            # self.img_id_to_dir = region_dicts["img_id_to_dir"]
            fd.close()
        
        self.vocab_size = len(region_dicts["vocab"])
        self.bos_idx = self.vocab_size + 1
        self.eos_idx = self.vocab_size + 1 
        self.eos_token = "&lt;eos&gt;"
        self.bos_token = "&lt;bos&gt;"
        self.pad_idx = self.vocab_size + 2
        vocab_size = self.pad_idx
        
        self.image_vector_dim = cfg.LANGUAGE_MODEL.FEATS_DIM  
        self.input_encoding_size = cfg.LANGUAGE_MODEL.FEATS_ENCODE_DIM
        self.rnn_type = cfg.LANGUAGE_MODEL.CELL_TYPE  
        self.rnn_size = cfg.LANGUAGE_MODEL.HIDDEN_DIM  
        self.num_layers = cfg.LANGUAGE_MODEL.NUM_LAYERS  
        self.dropout = 0.5  # opt.drop_prob_lm
        self.seq_length = cfg.LANGUAGE_MODEL.SEQUENCE_LEN
        
        # sample from the distribution instead
        self.sample_argmax = True
        
        sequence_core = LSTM(input_dim=self.input_encoding_size, hidden_dim=self.rnn_size)
        dropout = nn.Dropout(self.dropout) if self.dropout else None
        logit = nn.Linear(self.rnn_size, vocab_size + 1)

        rnn = nn.Sequential()
        rnn.add_module("sequence_core", sequence_core)
        rnn.add_module("dropout", dropout)
        rnn.add_module("logit", logit)
        
        input_embedding = ParallelTable(image_input_dim=self.image_vector_dim,
                                        image_embed_dim=self.input_encoding_size,
                                        rnn_size=self.rnn_size, vocab_size=vocab_size)
        self.net = nn.Sequential()
        self.net.add_module("input_embedding", input_embedding)
        # self.net.add_module("rnn", nn.Sequential(*rnn_block))
        self.net.add_module("rnn", rnn)
    
    def pad_sequence(self, gt_sequence):
        N, T = gt_sequence.shape
        gt_pad = gt_sequence.new_zeros(N, T + 1)
        gt_pad[:, 0] = self.bos_idx
        gt_pad[:, 1:T + 1] = gt_sequence
        mask = torch.eq(gt_pad, 0)
        gt_pad[mask] = self.pad_idx
        return gt_pad
    
    def get_gt_caps(self, gt_sequence, gt_lengths):
        N, T = gt_sequence.shape
        if gt_sequence.numel() == 0:
            print(gt_sequence)
        assert gt_sequence.numel() &gt; 0, "panic, failed ground truth"
        target = torch.zeros(N, T + 2, dtype=gt_sequence.dtype)
        target[:, 1:T + 1] = gt_sequence
        target[torch.arange(N), gt_lengths + 1] = self.eos_idx
        return target.to(gt_sequence.device)
    
    def _reset_states(self):
        for i in range(self.num_layers):
            rnn_size = len(self.net.rnn)
            for r in range(rnn_size):
                layer = self.net.rnn[r]
                if type(layer) == LSTM:
                    layer.reset_states()
                    layer.remember_states = True
                    break
    
    def setup_optimizer(self, model):
        import torch.optim
        params = []
        BASE_LR = 0.0005
        WEIGHT_DECAY = 0.0005
        WEIGHT_DECAY_BIAS = 0.0
        BIAS_LR_FACTOR = 2
        BETAS = [0.9, 0.999]
        EPS = 1e-8
        
        for key, value in model.named_parameters():
            if not value.requires_grad:
                print("%s donot need grad" % key)
                continue
            
            lr = BASE_LR
            weight_decay = WEIGHT_DECAY
            if "bias" in key:
                lr = BASE_LR * BIAS_LR_FACTOR
                weight_decay = WEIGHT_DECAY_BIAS
            params += [{"params": [value], "lr": lr, "weight_decay": weight_decay}]
        optimizer = torch.optim.Adam(
            params,
            lr=BASE_LR,
            betas=BETAS,
            eps=EPS
        )
        
        for p in self.state_dict():
            print(p)

        return optimizer
    
    def forward(self, feats, seqs):
        pad_seqs = self.pad_sequence(seqs)
        ret = self.net((feats, pad_seqs))
        return ret
    
    def sample(self, fc_feats):
       # sample_method = opt.get('sample_method', 'greedy')
        
        batch_size = fc_feats.size(0)
        seq = fc_feats.new_zeros(batch_size, self.seq_length, dtype=torch.long)
        # seqLogprobs = fc_feats.new_zeros(batch_size, self.seq_length)
        self._reset_states()
        
        # First timestep: image vectors, ignore output
        image_vecs_encoded = self.net.input_embedding.image_encoder.forward(fc_feats)
        # noised_img = image_vecs_encoded + torch.randn(image_vecs_encoded.shape[0], 512).to("cuda")
        self.net.rnn.forward(image_vecs_encoded.unsqueeze(1))
        
        for t in range(self.seq_length):
            if t == 0:
                words = fc_feats.data.new(batch_size).long().zero_().fill_(self.bos_idx)
            else:
                words = seq[:, t - 1]
            
            wordvecs = self.net.input_embedding.lookup_table.forward(words)
            scores = self.net.rnn.forward(wordvecs.unsqueeze(1)).view(batch_size, -1)
            if self.sample_argmax:
                _, idx = torch.max(scores, 1)
                idx = idx.view(-1).long()
            else:
                logprobs = F.softmax(self.word_logit(scores), dim=1)
                idx = torch.multinomial(logprobs, 1)
                idx = idx.view(-1).long()
            seq[:, t] = idx
        print(seq)
        return seq


if __name__ == "__main__":
    ################################################################################

    loss_fn = TemporalCrossEntropyCriterion()
    toy_model = RNN_LanguageModel(cfg=cfg).to("cuda")
    
    optim = toy_model.setup_optimizer(lm)
    optim.zero_grad()
    for iter in range(1000):
        pred = toy_model.forward(feats=train_feats, seqs=train_seqs)
        
        loss = loss_fn(pred, toy_model.get_gt_caps(train_seqs, train_lens))
        loss.backward()
        if iter % 10 == 0:
              print("Iteration: {} loss: {}".format(iter, loss))
        optim.step()
        optim.zero_grad()
    ################################################################################
    
    toy_model.eval()
    pred = toy_model.sample(test_feats)
</code></pre>
<p>It might be helpful to directly have a look the model structure</p>
<pre><code class="lang-auto">RNN_LanguageModel(
  (net): Sequential(
    (input_embedding): ParallelTable(
      (image_encoder): Sequential(
        (0): Linear(in_features=4096, out_features=512, bias=True)
        (1): ReLU(inplace=True)
      )
      (lookup_table): Embedding(910, 512)
    )
    (rnn): Sequential(
      (sequence_core): LSTM()
      (dropout): Dropout(p=0.5, inplace=False)
      (logit): Linear(in_features=512, out_features=909, bias=True)
    )
  )
)

</code></pre>
<p>Here is the training log:</p>
<pre><code class="lang-auto">Iteration: 0 loss: 18.254560470581055
Iteration: 10 loss: 17.90742301940918
Iteration: 20 loss: 17.980499267578125
Iteration: 30 loss: 17.781349182128906
Iteration: 40 loss: 17.84979248046875
Iteration: 50 loss: 17.896930694580078
Iteration: 60 loss: 17.904727935791016
Iteration: 70 loss: 17.801372528076172
Iteration: 80 loss: 18.06501579284668
Iteration: 90 loss: 17.91199493408203
Iteration: 100 loss: 17.90298843383789
Iteration: 110 loss: 17.859905242919922
Iteration: 120 loss: 18.082622528076172
Iteration: 130 loss: 17.864147186279297
Iteration: 140 loss: 18.0921630859375
.....
</code></pre>
<p>Seems  like it does not intend to converge. The Idea is based on paper <a href="https://github.com/jcjohnson/torch-rnn/blob/master/LSTM.lua" rel="noopener nofollow ugc">DenseCap: Fully Convolutional Localization Networks for Dense Captioning</a></p>