Hi I am using LSTMcell as the decoder to complete a task similar to machine translation. However, I met the following GPU out of memory problem:
RuntimeError: CUDA out of memory. Tried to allocate 2.00 MiB (GPU 0; 23.64 GiB total capacity; 18.54 GiB already allocated; 8.06 MiB free; 22.44 GiB reserved in total by PyTorch)
I checked the gpu memory usage during the running time of the decoder with print(f"allcoated {torch.cuda.memory_allocated(0)/(1024**3)} GB") for the my customized attention module, and saw the memory usage keeps climbing.
Detaild error message points to the Attention module:
<ipython-input-15-ddbbf84c3dca> in forward(self, encoder_out, decoder_hidden)
29 print("---")
30 # alpha: calculated base level attention
---> 31 alpha = self.softmax(att) # (batch, sequence_len); row sum = 1
32 attention_weighted_encoding = encoder_out * alpha.unsqueeze(2)
33 attention_weighted_encoding = torch.sum(attention_weighted_encoding, dim=1) # (batch, encoder_dim)
~/anaconda3/envs/torch_dev/lib/python3.8/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
725 result = self._slow_forward(*input, **kwargs)
726 else:
--> 727 result = self.forward(*input, **kwargs)
728 for hook in itertools.chain(
729 _global_forward_hooks.values(),
~/anaconda3/envs/torch_dev/lib/python3.8/site-packages/torch/nn/modules/activation.py in forward(self, input)
1196
1197 def forward(self, input: Tensor) -> Tensor:
-> 1198 return F.softmax(input, self.dim, _stacklevel=5)
1199
1200 def extra_repr(self) -> str:
~/anaconda3/envs/torch_dev/lib/python3.8/site-packages/torch/nn/functional.py in softmax(input, dim, _stacklevel, dtype)
1510 dim = _get_softmax_dim('softmax', input.dim(), _stacklevel)
1511 if dtype is None:
-> 1512 ret = input.softmax(dim)
1513 else:
1514 ret = input.softmax(dim, dtype=dtype)
Here is the code of the attention module:
class Attention(nn.Module):
"""
Attention network for calculate attention value
"""
def __init__(self, encoder_dim, decoder_dim, attention_dim):
"""
:param encoder_dim: input size of encoder network
:param decoder_dim: input size of decoder network
:param attention_dim: input size of attention network
"""
super(Attention, self).__init__()
self.encoder_att = nn.Linear(encoder_dim, attention_dim)
self.decoder_att = nn.Linear(decoder_dim, attention_dim)
self.full_att = nn.Linear(attention_dim, 1)
self.relu = nn.ReLU()
self.softmax = nn.Softmax(dim=1)
def forward(self, encoder_out, decoder_hidden):
att1 = self.encoder_att(encoder_out) # (batch, sequence_len, encoder_dim) -> (batch, sequence_len, attention_dim)
att2 = self.decoder_att(decoder_hidden) # (batch, decoder_dim) -> (batch, attention_dim)
print(att1.shape, att2.shape)
print(f"allcoated {torch.cuda.memory_allocated(0)/(1024**3)} GB")
# base level attention
att = self.full_att(self.relu(att1 + att2.unsqueeze(1))).squeeze(2) # (batch, sequence_len)
print(f"allcoated {torch.cuda.memory_allocated(0)/(1024**3)} GB")
print("---")
# alpha: calculated base level attention
alpha = self.softmax(att) # (batch, sequence_len); row sum = 1
attention_weighted_encoding = encoder_out * alpha.unsqueeze(2)
attention_weighted_encoding = torch.sum(attention_weighted_encoding, dim=1) # (batch, encoder_dim)
return attention_weighted_encoding, alpha
And here is how decoder works in my code:
# predict sequence
for t in range(max(decode_lengths)):
batch_size_t = sum([l > t for l in decode_lengths])
# attention_weighted_encoding: (batch_size_t, encoder_dim)
# alpha: (batch_size_t, num_of_pixels)
attention_weighted_encoding, alpha = self.attention(encoder_out[:batch_size_t], h[:batch_size_t])
gate = self.sigmoid(self.f_beta(h[:batch_size_t])) # gating scalar, (batch_size_t, encoder_dim)
attention_weighted_encoding = gate * attention_weighted_encoding
# new hidden and new cell state
h, c = self.decode_step(
torch.cat([embeddings[:batch_size_t, t, :], attention_weighted_encoding], dim=1),
(h[:batch_size_t], c[:batch_size_t])
)
preds = self.fc(self.dropout(h)) # (batch_size_t, vocab_size)
predictions[:batch_size_t, t, :] = preds
alphas[:batch_size_t, t, :] = alpha
I appreciate any input that could help me solve this issue. Thanks in advance!