[Require Help] Pytorch: Saving Tensor to List result in CUDA Out of Memory

username: Daniel_Kurniadi
cooked: <h2>[Require Help] Pytorch: Saving Tensor to List result in CUDA Out of Memory</h2>
<p>I’m currently experimenting with cascaded model for video denoising and video action classifier for my experimentation. Anyway, I encountered an <code>CUDA Runtime Out of Memory Error: </code> which I have investigated and turns out I need a better way to optimise my code, but Idk how. I need help regarding this matter. Here, let me go through my code and see if you guys have any idea how to fix this better.</p>
<p><em>Note</em>: My models are rigid in the input shape for sequence/clip length (time) and image size.</p>
<p>Before, there are some terminology I use:</p>
<ul>
<li>
<code>[N, T, C, H, W]</code> - in this order, is input shape of data coming from <code>DALIVideoLoader</code>, which is my data loader. <code>N</code> is <code>batch_size</code>, <code>T</code> is the clip length per iteration, and <code>[C, H, W]</code> is the RGB frame shape</li>
<li>I have a sliding window for the time dimension, denoted as small <code>t</code>. for each window of time it will contain <code>[N, t, C, H, W]</code> of data where <code>t &lt; T</code>.</li>
<li>
<code>FastDVDNet.NUM_INPUT_FRAMES</code> is the input to denoiser model, I will explain below</li>
<li>
<code>sequence_length</code> - the clip length, denoted by <code>T</code> above</li>
</ul>
<h3>Flow, part of question</h3>
<ol>
<li>
<code>STEP1</code>: Suppose I receive <code>[8, 20, 3, 224, 224]</code> (5 dimension) tensor raw batch of data from my data loader, I need to perform sliding window w.r.t time with window_size <code>t=5</code> over the time dimension. Hence each slide, I extract the clip of shape <code>[8, 5, 3, 224, 224]</code> which in total I will only get 16 sequences of the mentioned shape</li>
<li>
<code>STEP2</code>: Now for each sequence I get above, I feed to the denoiser model. Denoiser model receive <code>[8, 5, 3, 224, 224]</code> and output map to <code>[8, 3, 224, 224]</code>. Now I <strong>collect those <code>16 x [8, 3, 224, 224]</code>  denoised tensors in a python list</strong>.</li>
<li>
<code>STEP3</code>: I finally stack the tensors in the python list mentioned into a single 5-dimension <code>[8, 16, 3, 224, 224]</code> shaped tensor. Feed this to the classifier model and it output me a prediction.</li>
</ol>
<h3>Problem and Question:</h3>
<p>As mentioned, I got <code>CUDA Out of Memory error</code>. I know I can add more GPU/RAM but I wonder if there is anything I can fix from my code to reduce the overhead. Currently I observe that <strong>collecting<br>
the tensors in a list</strong> is the cause of exploding memory usage (up to <code>10GB</code> despite I have reduced the batch_size <code>N</code> to 1). But how else can I do to perform the above Flow (see section above)? Please kindly advice.</p>
<h3>Code:</h3>
<p>Here is a sample code if you wanted to take a look and for your reference.</p>
<pre><code class="lang-auto">def main(**args):
	""" Performs the main training loop
	"""
	# Load dataset
	print('&gt; Loading datasets ...')

	train_loader = DaliVideoLoader(file_list = args['trainset_list'],
									batch_size = args['batch_size'],
									sequence_length = args['temp_patch_size'],
									step = 1)

	# STEP 0: Precalculate input-output sizes
	sequence_length = args['temp_patch_size']
	t = window_size = FastDVDnet.NUM_INPUT_FRAMES  # here, t=5
	den_stride = 1
	den_frame_length = int(sequence_length - window_size + den_stride) // den_stride

	# Model definition
	den_model = FastDVDnet().cuda()
	classifier_model = MultiFiberNet3d(num_classes=51)

	for data, labels in train_loader:  # CUDA torch.Tensor loader
		# shape of data: [N, T, C, H, W], e.g. T = 20
		# shape of label: [N, ]

		# STEP 1: Partition one video clip into segments of t=5 frames.
		# sliding window to get tuple of tensor: [N, t=5, C, H, W] x den_frame_length
		# e.g. t = 5
		frame_patches = tuple(data[:, m:m+t, :, :, :] for m in range(0, den_frame_length))

		print("FRAME", len(frame_patches), frame_patches[-1].shape)

		den_frames_stack = []

		# STEP 2: Sample noise and Denoise Frame with video denoiser model 
		for frames in frame_patches:
			# frame is of shape [N, t=5, C, H, W]

			# std dev of each sequence
			stdn = torch.rand(N, 1, 1, 1, 1).cuda()
			eps = torch.randn_like(stdn)
			# draw noise samples from std dev tensor
			noise = eps.mul(stdn)

			frames = frames + noise

			frames = frames.cuda(non_blocking=True)
			noise = noise.cuda(non_blocking=True)

			noise_map = stdn.expand((N, 1, 1, H, W)).cuda(non_blocking=True) # one channel per image

			print("NOISE MAP", noise_map.shape)

			# Evaluate model and optimize it
			den_frames = den_model(frames, noise_map)

			# Then put to list and BOOM, it OOM after 10-20 batches!
			den_framden_frames_stackes.append(den_frames)  
			print("DEN FRAME", den_frames.shape)

		# STEP 3: Combine all denoised frames into one tensor 
		# frames_train is of shape: [N, F, C, H, W]
		# where F is calculated as `den_frame_length` above
		frames_train = torch.stack(den_frames_stack, dim=1)

		output = classifier_model(frames_train)

	# Print elapsed time
	elapsed_time = time.time() - start_time
	print('Elapsed time {}'.format(time.strftime("%H:%M:%S", time.gmtime(elapsed_time))))
</code></pre>

username: ptrblck
cooked: <p>I don’t see any obvious way to lower the memory usage.<br>
If you’ve already tried to use a batch size of 1, you might need to use the <a href="https://pytorch.org/docs/stable/checkpoint.html">checkpoint</a> utility to trade compute for memory.</p>