Trying to train LSTM on GPU

username: Felix_Labelle
cooked: <p>I’ve been trying to train an LSTM cell using a GPU, which works on a CPU, but get the following error.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/2X/2/23f7ae46774d12f9208174a3a49ab812dfab1085.png" data-download-href="https://discuss.pytorch.org/uploads/default/23f7ae46774d12f9208174a3a49ab812dfab1085" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/2X/2/23f7ae46774d12f9208174a3a49ab812dfab1085.png" alt="image" data-base62-sha1="58bmDgexmJKFaHKu5NiQkls2Yex" width="690" height="183" data-dominant-color="0B0B0B"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1010×268 10 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>I’ve found this post (<a href="https://discuss.pytorch.org/t/how-to-train-lstm-with-gpu/32466/17" class="inline-onebox">How to train LSTM with GPU - #17 by Crohn_eng</a>), but have been using a custom collate function and haven’t found the answer to my issue in this post.</p>
<p>Below is the LSTM code:</p>
<pre><code class="lang-auto">import torch
import torch.nn.utils.rnn as rnn_utils
import torch.nn as nn
from torchUtils import SplitDataset

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
def my_collate(batch):
    # batch contains a list of tuples of structure (sequence, target)
    data = [item[0] for item in batch]
    data = rnn_utils.pack_sequence(data, enforce_sorted=False)
    targets = [item[1] for item in batch]
    return [data, targets]
	

dataset = SplitDataset(shuffleData=False) #ImageDataset()
dataset.selectPhase('train')
trainingDL = torch.utils.data.DataLoader(dataset, 4,shuffle=True, num_workers=0, collate_fn=my_collate)
inputs, labels = next(iter(trainingDL))
print(inputs[0].shape)

'''
a = torch.Tensor([[1, 2], [2, 3], [3, 4]])
b = torch.Tensor([[4, 5], [5, 6]])
c = torch.Tensor([[6, 7]])
packed = rnn_utils.pack_sequence([a, b, c])
'''

lstm = nn.LSTM(1000,1000, batch_first=True).to(device)
#device = torch.device("cpu")

packed_output, (h,c) = lstm(inputs.to(device))

y = rnn_utils.pad_packed_sequence(packed_output)
predictionLayer = nn.Linear(1000, 2)
print(len(y))
print(y[1])
print(y[0].shape)
currentIdx = -1
for count, pos in enumerate(y[1]):
	print(pos, count)
	print(predictionLayer(y[0][pos - 1, count, :]))
</code></pre>
<p>The custom dataset is the following:</p>
<pre><code class="lang-auto">class SplitDataset(MongoDataset):
	def __init__(self, trainingProportion = 0.8, shuffleData = False):
		MongoDataset.__init__(self)
		trainingSetSize = int(len(self)*trainingProportion)
		validSetSize =  len(self) - trainingSetSize
		# Could be replaced by something else rather than random replacement
		random.shuffle(self.ids)
		self.phase = 'train'
		self.trainingSet = self.ids[0:trainingSetSize]
		self.validationSet = self.ids[trainingSetSize:]
		self.shuffleData = shuffleData
		self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
		self.featureExtractor = models.vgg16(pretrained=True)

	def __getitem__(self, i):
		profile, label = MongoDataset.__getitem__(self,i)
		# You could replace this with the online data loader
		with torch.no_grad():
			self.featureExtractor.eval()
			preppedData = [self.featureExtractor(self.imageTransform(img).unsqueeze(0)).squeeze().to(self.device) for img in profile.images]

		# TODO: Replace this a proper transform using Pytorch libs
		if self.phase == 'train' and self.shuffleData:
			np.random.shuffle(preppedData)
		return torch.stack(preppedData), label
		
	def selectPhase(self, phase):
		self.phase = phase
		if phase == 'val':
			self.loadData = self.validationSet
			self.transform = ProfileDataset.validationTransform
		elif phase == 'train':
			self.loadData = self.trainingSet
			self.transform = ProfileDataset.trainingTransform
		else:
			raise("Invalid phase")
</code></pre>
<p>The error appears to be caused by the model’s weights, is it possible the model is only partially on the GPU?</p>

username: jerinphilip
cooked: <aside class="quote no-group" data-username="Felix_Labelle" data-post="1" data-topic="47674">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/felix_labelle/48/11625_2.png" class="avatar"> Felix_Labelle:</div>
<blockquote>
<p>y[0][pos - 1, count, :])</p>
</blockquote>
</aside>
<p><code>y[0][pos - 1, count, :]).to(device)</code></p>

username: Felix_Labelle
cooked: <p>That wasn’t the solution, but was helpful. It was a silly mistake, the prediction layer wasn’t sent to device. Adding</p>
<pre><code class="lang-auto">predictionLayer = nn.Linear(1000, 2).to(device)
</code></pre>
<p>fixed it.</p>