Hi, I implemented an attention-based Sequence-to-sequence model in Theano and then ported it into PyTorch.
However, the GPU memory usage in Theano is only around 2GB, while PyTorch requires almost 5GB, although it’s much faster than Theano.
Maybe it’s a trading consideration between memory and speed. But the GPU memory usage has increased by 2.5 times, that is unacceptable. I think there should be room for optimization to reduce GPU memory usage and maintaining high efficiency.
I printed out the allocated, max_allocated, cached and max_cached GPU memory after zero_grad, model forwarding, loss calculation, backward and optimizer step.
I don’t quite understand how memory usage changes according to the following output.
1). In the 0-th step, max allocated memory rapidly increased from 2154M to 2566M after “compute loss”. Is this normal?
2). In the 0-th step, after “backward”, max allocated memory increased from 2566M to 3482M again. The increased 916M should consists of Adam's exp_avg and exp_avg_sq for each parameter, so the expected increasement should be about 2*281M=562M instead of 916M. according to my understanding, all the intermediate results required to calculate derivatives are already stored in the forward pass, so there should be no more memory allocation during backward.
3). In the 0-th step, after “backward”, allocated memory is 1020M, while after optimizer step, the allocated memory is 1581M. What’s taking up the extra 1300M (1581M - 281M) here? Although I invoke model.zero_grad() in the 1-st step, the allocated memory is still 1581M, it seems that PyTorch didn’t release any GPU memory.
4). In the 1-th step, after “backward”, why max allocated memory increased 847M again, from 3482M to 4329M, as in 2) ? Note that during the training, each batch are totally the same.
5). Finally, is there anything in my code that can be optimized to reduce the GPU memory usage?
Thanks in advance. Really appreciate if anyone explain these problems for me.
Below is the training output and my code.
-- beginning
allocated: 281M, max allocated: 281M, cached: 284M, max cached: 284M
-- zero grads
allocated: 281M, max allocated: 281M, cached: 284M, max cached: 284M
-- forward
allocated: 2108M, max allocated: 2154M, cached: 2178M, max cached: 2178M
-- compute loss
allocated: 2566M, max allocated: 2566M, cached: 2635M, max cached: 2635M
-- backward
allocated: 1020M, max allocated: 3482M, cached: 3551M, max cached: 3551M
-- optimizer step
allocated: 1581M, max allocated: 3482M, cached: 3551M, max cached: 3551M
***************** finish train step: 0

-- zero grads
allocated: 1581M, max allocated: 3482M, cached: 3551M, max cached: 3551M
-- forward
allocated: 2955M, max allocated: 3482M, cached: 3552M, max cached: 3552M
-- compute loss
allocated: 3413M, max allocated: 3482M, cached: 3552M, max cached: 3552M
-- backward
allocated: 1581M, max allocated: 4329M, cached: 4468M, max cached: 4468M
-- optimizer step
allocated: 1581M, max allocated: 4329M, cached: 4468M, max cached: 4468M
***************** finish train step: 1

-- zero grads
allocated: 1581M, max allocated: 4329M, cached: 4468M, max cached: 4468M
-- forward
allocated: 2955M, max allocated: 4329M, cached: 4468M, max cached: 4468M
-- compute loss
allocated: 3413M, max allocated: 4329M, cached: 4468M, max cached: 4468M
-- backward
allocated: 1581M, max allocated: 4329M, cached: 4468M, max cached: 4468M
-- optimizer step
allocated: 1581M, max allocated: 4329M, cached: 4468M, max cached: 4468M
***************** finish train step: 2

Please refer to following code as a ready-to-run example of my model and training loop.
import argparse

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.utils.rnn import pad_packed_sequence as unpack, pack_padded_sequence as pack
from torch.optim.adam import Adam


class Encoder(nn.Module):
    def __init__(self, args):
        super().__init__()
        self.args = args

        self.rnn = nn.GRU(args.input_size, args.hidden_size, bidirectional=True, batch_first=True)
        self.embedding = nn.Embedding(args.vocab_size, args.input_size, padding_idx=args.pad_id)

    def forward(self, input, h0=None):
        """
        Args:
            input: long tensor of shape (B x T)
            h0:
        """
        mask = input.eq(self.args.pad_id)
        total_length = input.size(1)
        # Trim mask when trained with data-parallel
        lens = input.size(1) - mask.sum(1)
        if torch.cuda.device_count() > 1:
            max_len = lens.max().item()
            mask = mask[:, :max_len]
        x = self.embedding(input)
        x = pack(x, lens.tolist(), True)

        x, hx = self.rnn(x, h0)
        x, _ = unpack(x, batch_first=True, total_length=total_length)

        summary = hx[1]  # The last state of backward rnn
        return {
            'feature': x,
            'hidden': summary,
            'mask': mask
        }


class Bridge(nn.Module):
    r"""
    Bridge the connection between encoder and decoder.
    """

    def __init__(self, args):
        super().__init__()
        self.map_k = nn.Sequential(nn.Linear(args.hidden_size * 2, args.attention_size),
                                   nn.Tanh())
        self.dec_init = nn.Sequential(nn.Linear(args.hidden_size, args.hidden_size),
                                      nn.Tanh())

    def forward(self, encoder_state):
        K = self.map_k(encoder_state['feature'])
        h0 = self.dec_init(encoder_state['hidden'])
        return {
            'key': K,
            'value': encoder_state['feature'],
            'mask': encoder_state['mask'],
            'hidden': h0,
        }


class Attention(nn.Module):
    def __init__(self, query_size, attention_size):
        super(Attention, self).__init__()
        self.query_size = query_size
        self.key_size = attention_size
        self.map_query = nn.Linear(query_size, attention_size)
        self.v = nn.Linear(attention_size, 1)

    def forward(self, query, keys, values, mask):
        """
        score = v^T * (tanh(W * K + U * q))
        Args:
            query: B x D
            keys: B x T x D
            values: B x T x D
            mask: B x T
        """
        # B x T x D
        x = keys + self.map_query(query).unsqueeze(1)
        # B x T
        x = self.v(torch.tanh(x)).squeeze(-1)
        x.data.masked_fill_(mask, -float('inf'))
        x = F.softmax(x, -1)
        output = torch.bmm(x.unsqueeze(1), values).squeeze(1)
        return output, x


class Decoder(nn.Module):
    def __init__(self, args):
        super().__init__()

        input_size = args.input_size
        hidden_size = args.hidden_size
        ctx_size = 2 * hidden_size

        self.embedding = nn.Embedding(args.vocab_size, input_size, args.pad_id)

        self.cell0 = nn.GRUCell(input_size, hidden_size)
        self.attention = Attention(hidden_size, args.attention_size)
        self.cell1 = nn.GRUCell(ctx_size, hidden_size)

        readout = nn.Sequential(
            nn.Linear(input_size + hidden_size + ctx_size, input_size),
            nn.Tanh()
        )

        logit = nn.Linear(input_size, args.vocab_size)

        self.generator = nn.Sequential(
            readout,
            logit
        )

        if args.weight_tying:
            self.logit.weight = self.embedding.weight

    def forward(self, inputs, state):
        x = self.embedding(inputs)
        K, V, mask, h = state['key'], state['value'], state['mask'], state['hidden']
        hs = []
        ctxs = []
        for x_i in x.split(1, 1):
            x_i = x_i.squeeze(1)  # B x D
            h = self.cell0(x_i, h)
            ctx, _ = self.attention(h, K, V, mask)
            h = self.cell1(ctx, h)
            hs.append(h)
            ctxs.append(ctx)
        hs = torch.stack(hs, 1)
        ctxs = torch.stack(ctxs, 1)
        logit = self.generator(torch.cat([x, hs, ctxs], -1))
        state['hidden'] = h
        return logit, state


class Seq2Seq(nn.Module):
    def __init__(self, args):
        super().__init__()
        self.args = args
        self.encoder = Encoder(args)
        self.decoder = Decoder(args)
        self.bridge = Bridge(args)

    def encode(self, input):
        state = self.encoder(input)
        return state

    def decode(self, input, state):
        logit, state = self.decoder(input, state)
        return logit, state

    def forward(self, input_x, input_y):
        enc_state = self.encode(input_x)
        dec_state = self.bridge(enc_state)
        del enc_state
        # Reuse eos as bos to left-pad the target.
        input_y = F.pad(input_y, [1, 0], mode='constant', value=self.args.eos_id)
        # Exclude eos from input_y
        input_y = input_y[:, :-1]
        logit, dec_state = self.decode(input_y, dec_state)
        return logit


def stat_cuda(msg):
    print('--', msg)
    print('allocated: %dM, max allocated: %dM, cached: %dM, max cached: %dM' % (
        torch.cuda.memory_allocated() / 1024 / 1024,
        torch.cuda.max_memory_allocated() / 1024 / 1024,
        torch.cuda.memory_cached() / 1024 / 1024,
        torch.cuda.max_memory_cached() / 1024 / 1024
    ))


def fake_batch(n_batch, batch_size, max_length=50):
    for _ in range(n_batch):
        x = torch.ones(batch_size, max_length).long().cuda() * 2
        y = torch.ones(batch_size, max_length).long().cuda() * 2
        yield (x, y)


def main(args):
    model = Seq2Seq(args).cuda()
    criterion = nn.CrossEntropyLoss(ignore_index=args.pad_id, reduction='sum').cuda()
    optimizer = Adam(model.parameters())
    for p in model.parameters():
        nn.init.uniform_(p, -0.1, 0.1)

    stat_cuda('beginning')
    for i, batch in enumerate(fake_batch(100, args.batch_size)):
        x, y = batch
        batch_size, length = x.size()
        # 1. zero grads
        model.zero_grad()
        stat_cuda('zero grads')
        # 2. forward
        logit = model(x, y)
        stat_cuda('forward')
        # 3. compute loss
        logit = logit.view(-1, logit.size(-1))
        loss = criterion(logit, y.view(-1))
        loss = loss / batch_size
        stat_cuda('compute loss')
        # 4. backward
        loss.backward()
        stat_cuda('backward')
        # 5. optimizer step
        optimizer.step()
        stat_cuda('optimizer step')
        print('***************** finish train step: {}\n'.format(i))

if __name__ == '__main__':
    args = argparse.Namespace()
    args.attention_size = 1000
    args.hidden_size = 1000
    args.input_size = 500
    args.pad_id = 0
    args.vocab_size = 30000
    args.weight_tying = False
    args.batch_size = 80
    args.eos_id = 1
    main(args)