Hi folks,
I’m using wavenet to predict n time steps base on a sequence. The training code is like this.
def train_batch(self, input_batch, target_batch):
    input_batch = Variable(input_batch)
    target_batch = Variable(target_batch)

    # Zero gradients of both optimizers
    self.optimizer.zero_grad()

    target_len = target_batch.size(0)
    batch_size = input_batch.size(0)

    # Move new Variables to CUDA
    predictions = Variable(torch.zeros(target_len, batch_size))
    if USE_CUDA:
        input_batch = input_batch.cuda()
        target_batch = target_batch.cuda()
        predictions = predictions.cuda()

    for t in range(target_len):
        output_batch = wavenet(input_batch)
        predictions[t] = output_batch[:, 0, -1]
        input_batch = torch.cat([input_batch[:, :, 1:], output_batch[:, :, -1], dim=2)

    # Loss calculation and backpropagation
    loss = self.criterion(predictions, target_batch)
    loss.backward()

I found that the GPU memory is proportional to the target_len. My understanding is pytorch allocates memory for wavenet instance for every time step. If I’m right, it is not necessary because calculating each time step needs the output of last time step makes the process actually executed sequentially, hence one wavenet instance is enough. The result is my implementation which contains one wavenet layer takes about 5G GPU. Another tensorflow implementation which contains three wavenet layers can run in a 8G GPU. Is there any solution? Is there any way to control the allocation of the memory?
Thanks!
BTW, a similar topic is here:

  
    
    
    GPU memory usage on long seq2seq sequences 
  
  
    Something I really don’t understand is whether RNN are unrolled or not and how to choose. In Keras with Theano backend you can specify unroll = True or False depending on the memory vs speed tradeoff you want to make. 
By the way in the specific case described here, whatever the rolling status of the LSTM, the memory consumption will grow due to the output you use: 
net_out = nn.Linear(in_features=HIDDEN_SIZE, out_features=INPUT_SIZE) 
Indeed : 
print(net_out) 
Gives : 
Linear (512 -> INPUT_SIZE…
  


The difference is they found the GPU memory is proportional to the input length and I found it is also proportional to the output length.