RAM out of memory and process killed from 5 epoch

MyeongJin: <p>My RAM usage keeps on increasing after first epoch.<br>
The entire time GPU memory remains constant. Only RAM increases.<br>
RAM increases 1 to 2% per an iteration, and after 4 epochs, process killed .<br>
I can’t find what is the problem in my codes.</p>
<pre><code class="lang-auto">#coding:utf8
import ipdb;
import os
import models
import torch as t
import time
import scipy.io as sio
import numpy as np
import cv2
from torch.utils.tensorboard import SummaryWriter

from config import opt
from data.dataset import *
from torch.utils.data import DataLoader
from tqdm import tqdm
from torchnet import meter
from skimage.metrics import peak_signal_noise_ratio
from skimage.metrics import structural_similarity as ssim
import matplotlib.pyplot as plt

import scipy.io
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

@t.no_grad() # pytorch&gt;=0.5

def test(**kwargs):
   opt._parse(kwargs)
   
   writer = SummaryWriter(f'runs/{opt.env}_test_{time.strftime("%Y%m%d-%H%M%S")}')
   model = getattr(models, opt.model)(8 * 2 * opt.slice_num, 8 * 2 * opt.slice_num, 64, 5, 0, opt.slice_num, False)
   if opt.test_model_path:
       model.load(opt.test_model_path)
   model.to(opt.device)

   # data
   mask_smp_path = opt.undersample_mask #"./data/mask_smp.mat"
   test_data = Mridata(opt.val_data_root, mask_smp_path, slice_num=opt.slice_num, test=True)
   test_dataloader = DataLoader(test_data,opt.val_batch_size,
                       shuffle=False,num_workers=opt.num_workers,pin_memory=True)
   
   model.eval()
   rec_total_nmse  = 0
   rec_total_psnr  = 0
   rec_total_ssim  = 0
   data_total_nmse = 0
   data_total_psnr = 0
   data_total_ssim = 0
   
   result_dir = f'results_{opt.env}_{time.strftime("%Y%m%d-%H%M%S")}'
   os.makedirs(result_dir, exist_ok=True)
   
   heatmap_dir = os.path.join(result_dir,'heatmaps')
   result_img_dir = os.path.join(result_dir,'result_images')
   os.makedirs(heatmap_dir, exist_ok=True)
   os.makedirs(result_img_dir,exist_ok=True)
   
   for ii, (val_input, label, sen, dc_fft) in tqdm(enumerate(test_dataloader)):
       print(ii)
       val_input = val_input.to(opt.device)
       
       kl = fft2(label).data.numpy()   
       label = label.data.numpy() # b*c*W*h*2 real
       
       s_t = time.clock()
       val_pred, ks, ke, c_weight, c_weight2 = model(val_input)
       e_t = time.clock()
       print(f"Inference time: {e_t - s_t}")
       
       fig_heatmap, axs_heatmap = plt.subplots(1,11,figsize=(30,4))
       fig_heatmap.suptitle('Weight Heatmaps', fontsize=20)
       
       # weight map visualize
       for i in range(11):
           weight = c_weight[i][0,0,:,:].cpu().numpy()
           weight_normalized = (weight - weight.min()) / (weight.max() - weight.min())
           axs_heatmap[i].imshow(weight_normalized,cmap='jet')
           axs_heatmap[i].axis('off')
           axs_heatmap[i].set_title(f'Layer {i}', fontsize=12)
       
       plt.tight_layout()
       plt.subplots_adjust(wspace=0.05)
       plt.savefig(os.path.join(heatmap_dir,f'weight_heatmaps_{ii}.png'),dpi=300, bbox_inches='tight')
       plt.close(fig_heatmap)
           
       data_input = val_input.data.cpu().numpy()       
       val_pred = val_pred.data.cpu().numpy()
       sen = sen.data.numpy()       
       ks = ks.data.cpu().numpy() 
       ke = ke.data.cpu().numpy() 

       label = label[:,8:16,:,:,:]
       data_input = data_input[:,8:16,:,:,:]
       val_pred = val_pred[:,:,:,:,:]
       sen = sen[:,8:16,:,:,:]

       kl = kl[:,8:16,:,:,:]
       ks = ks[:,8:16,:,:,:]
       ke = ke[:,:,:,:,:]

       sen = sen[...,0] - 1j * sen[...,1]       #bcwh imag

       label = label[...,0] + 1j * label[...,1] #bcwh imag
       sen_label = np.sum(label * sen, 1)

       val_pred = val_pred[...,0] + 1j * val_pred[...,1] #bcwh imag
       sen_val_pred = np.sum(val_pred * sen, 1)

       data_input = data_input[...,0] + 1j * data_input[...,1]
       sen_data_input = np.sum(data_input * sen, 1)

       ks = ks[...,0] + 1j * ks[...,1]
       ks = np.sum(ks, 1)

       ke = ke[...,0] + 1j * ke[...,1]
       ke = np.sum(ke, 1)

       kl = kl[...,0] + 1j * kl[...,1]
       kl = np.sum(kl, 1)

       norval = np.absolute(sen_label.max())
       sen_label = np.absolute(sen_label) / norval
       sen_val_pred = np.absolute(sen_val_pred) / norval
       sen_data_input = np.absolute(sen_data_input) / norval
       ks = np.absolute(ks) / norval         
       ke = np.absolute(ke) / norval         
       kl = np.absolute(kl) / norval   
           
       # result plot
       fig, axs = plt.subplots(2,3,figsize=(20,10))

       # Label
       axs[0,0].imshow(sen_label.reshape(320,256),cmap='gray',vmax=0.03)
       axs[0,0].set_title('K-Input')
       axs[0,0].axis('off')

       # Input
       axs[0,1].imshow(sen_data_input.reshape(320,256),cmap='gray',vmax=0.03)
       axs[0,1].set_title('K-Output')
       axs[0,1].axis('off')

       # Output
       axs[0,2].imshow(sen_val_pred.reshape(320,256),cmap='gray')
       axs[0,2].set_title('K-Label')
       axs[0,2].axis('off')

       # K-input
       axs[1,0].imshow(ks.reshape(320,256),cmap='gray')
       axs[1,0].set_title('Input')
       axs[1,0].axis('off')

       # K-Output
       axs[1,1].imshow(ke.reshape(320,256),cmap='gray')
       axs[1,1].set_title('Output')
       axs[1,1].axis('off')

       # K-label
       axs[1,2].imshow(kl.reshape(320,256),cmap='gray')
       axs[1,2].set_title('Label')
       axs[1,2].axis('off')

       plt.tight_layout()
       plt.savefig(os.path.join(result_img_dir,f'result_image_{ii}.png'),dpi=300,bbox_inches='tight')
       plt.close(fig)
       
       rec_nmse = b_nmse(sen_label, sen_val_pred)
       rec_psnr = b_psnr(sen_label, sen_val_pred)
       rec_ssim = b_ssim(sen_label, sen_val_pred)

       data_nmse = b_nmse(sen_label, sen_data_input)
       data_psnr = b_psnr(sen_label, sen_data_input)
       data_ssim = b_ssim(sen_label, sen_data_input)
       
       writer.add_scalar('Metrics/REC_NMSE',rec_nmse,ii)
       writer.add_scalar('Metrics/REC_PSNR',rec_psnr,ii)
       writer.add_scalar('Metrics/REC_SSIM',rec_ssim,ii)
       writer.add_scalar('Metrics/DATA_NMSE',data_nmse,ii)
       writer.add_scalar('Metrics/DATA_PSNR',data_psnr,ii)
       writer.add_scalar('Metrics/DATA_SSIM',data_ssim,ii)
       
       rec_total_nmse += rec_nmse
       rec_total_psnr += rec_psnr
       rec_total_ssim += rec_ssim
       
       data_total_nmse += data_nmse
       data_total_psnr += data_psnr
       data_total_ssim += data_ssim

   rec_total_nmse /= opt.test_num 
   rec_total_psnr /= opt.test_num
   rec_total_ssim /= opt.test_num
   data_total_nmse /=opt.test_num
   data_total_psnr /=opt.test_num
   data_total_ssim /=opt.test_num
   
   writer.add_scalar('Final_Metrics/REC_NMSE',rec_total_nmse,0)
   writer.add_scalar('Final_Metrics/REC_PSNR',rec_total_psnr,0)
   writer.add_scalar('Final_Metrics/REC_SSIM',rec_total_ssim,0)
   writer.add_scalar('Final_Metrics/DATA_NMSE',data_total_nmse,0)
   writer.add_scalar('Final_Metrics/DATA_PSNR',data_total_psnr,0)
   writer.add_scalar('Final_Metrics/DATA_SSIM',data_total_ssim,0)

   print("REC  nmse: {} psnr: {} ssim: {}".format(rec_total_nmse, rec_total_psnr, rec_total_ssim))
   print("DATA nmse: {} psnr: {} ssim: {}".format(data_total_nmse, data_total_psnr, data_total_ssim))
   
   writer.close()
   
   with open(os.path.join(result_dir,'final_results.txt'),'w') as f:
       f.write("REC nmse: {} psnr: {} ssim: {}\n".format(rec_total_nmse,rec_total_psnr,rec_total_ssim))
       f.write("DATA nmse: {} psnr: {} ssim: {}\n".format(data_total_nmse,data_total_psnr,data_total_ssim))

def write_csv(results,file_name):
   import csv
   with open(file_name,'w') as f:
       writer = csv.writer(f)
       writer.writerow(['id','label'])
       writer.writerows(results)

def clip_gradient(optimizer, grad_clip):
   for group in optimizer.param_groups:
       for param in group["params"]:
           if param.grad is not None:
               param.grad.data.clamp(-grad_clip, grad_clip)

def weight_init(m):
   if isinstance(m, torch.nn.Conv2d):
       torch.nn.init.kaiming_uniform(m.weight)


def train(**kwargs):
   opt._parse(kwargs)

   writer = SummaryWriter(f'runs/{opt.env}_{time.strftime("%Y%m%d-%H%M%S")}')
   
   # step1: configure model
   model = getattr(models, opt.model)(8 * 2 * opt.slice_num, 8 * 2 * opt.slice_num, 64, 5, 0, opt.slice_num, False)
   if opt.load_model_path:
       model.load(opt.load_model_path)
   else:
       print('Initialize the model!')
       model.apply(weight_init)
   model.to(opt.device)

   # step2: data
   mask_smp_path = opt.undersample_mask #"./data/mask_smp.mat"
   
   train_data = Mridata(opt.train_data_root, mask_smp_path, slice_num=opt.slice_num, test=False)
   val_data = Mridata(opt.val_data_root, mask_smp_path, slice_num=1, test=True)
   train_dataloader = DataLoader(train_data,opt.batch_size,
                       shuffle=True,num_workers=opt.num_workers,pin_memory=True)
   val_dataloader = DataLoader(val_data,opt.val_batch_size,
                       shuffle=False,num_workers=opt.num_workers,pin_memory=True)
   
   # step3: criterion and optimizer
   criterion = t.nn.MSELoss(reduction='sum')
   lr = opt.lr
   optimizer = model.get_optimizer(lr, opt.weight_decay)
       
   # step4: meters
   loss_meter = meter.AverageValueMeter()
   #previous_loss = 1e10
   acc_n = 0
   acc_ave = 0
   
   # train
   for epoch in range(opt.max_epoch):
       epoch_loss =0 
       epoch_acc_ave = 0
       for ii,(data,label,sen,dc_fft) in tqdm(enumerate(train_dataloader)):
           # train model 
           input = data.to(opt.device)
           target = label.to(opt.device)
           optimizer.zero_grad()
           score, ks, ke, c_weight, c_weight2 = model(input)
           
           kl = fft2(label).cuda()
           ke_ = ke
           
           loss = criterion(score,target[:,8:16,:,:,:]) # + 10 * criterion(kl, ke_)  3 slice
           loss.backward()
           optimizer.step()
           
           loss_meter.add(loss.item())
           acc_m = acc_n + 320 * 256 * opt.batch_size
           acc_ave = (acc_n * acc_ave + loss) / acc_m     
           acc_n = acc_m

           print('Epoch: {} Iter: {} Loss: {} Objective: {} Env: {}'.format(epoch, ii, loss, acc_ave, opt.env))
           
           global_step = epoch * len(train_dataloader) + ii
           writer.add_scalar('Loss/train',loss.item(),global_step)
           writer.add_scalar('Objective',acc_ave.item(), global_step)
           if (ii+1) % opt.print_freq ==0:
               if os.path.exists(opt.debug_file):
                   ipdb.set_trace()
           
       writer.add_scalar('Loss/Train', epoch_loss / len(train_dataloader),epoch)
       writer.add_scalar('Objective', epoch_acc_ave, epoch)
       
       if (epoch) % 5 == 0:
           model.save('./checkpoints/' + opt.env + '_' + str(epoch) + '.pth')
       
       for param_group in optimizer.param_groups:
           if epoch &lt; 25:
               param_group['lr'] = opt.lr_list[epoch]
           else:
               param_group['lr'] = opt.lr_list[299]
       writer.add_scalar('Learning Rate',optimizer.param_groups[0]['lr'],epoch)
       if epoch % 1 ==0:
           KL = np.absolute(fft2(label).data.cpu().numpy()[0][...,0] + 1j*fft2(label).data.cpu().numpy()[0][...,1])[4,np.newaxis,:,:]
           KS = np.absolute(ks.data.cpu().numpy()[0][...,0] + 1j * ks.data.cpu().numpy()[0][...,1])[4,np.newaxis,:,:]
           KE = np.absolute(ke.data.cpu().numpy()[0][...,0] + 1j * ke.data.cpu().numpy()[0][...,1])[4,np.newaxis,:,:]
               
           writer.add_image('Label_kspace',KL[0],global_step,dataformats='HW')
           writer.add_image('Input_kspace',KS[0],global_step,dataformats='HW')
           writer.add_image('Output_kspace',KE[0],global_step,dataformats='HW')
           writer.add_image('Diff_kspace',(KS[0] - KE[0]),global_step,dataformats='HW')
               
           label_img = np.absolute(label.data.cpu().numpy()[0][...,0] + 1j*label.data.cpu().numpy()[0][...,1])[4]
           input_img = np.absolute(data.data.cpu().numpy()[0][...,0] + 1j*data.data.cpu().numpy()[0][...,1])[4]
           output_img = np.absolute(score.data.cpu().numpy()[0][...,0] + 1j*score.data.cpu().numpy()[0][...,1])[4]
               
           writer.add_image('Label',label_img,global_step,dataformats='HW')
           writer.add_image('Input',input_img,global_step,dataformats='HW')
           writer.add_image('Output',output_img,global_step,dataformats='HW')
           
       previous_loss = loss_meter.value()[0]
       
       model.zero_grad()
       torch.cuda.empty_cache()
   writer.close()
   
@t.no_grad()
def mse(gt, pred):
   """ Compute Mean Squared Error (MSE) """
   return np.mean((gt - pred) ** 2)

def b_nmse(gt, pred):
   """ Compute Normalized Mean Squared Error (NMSE) """
   batch_size = gt.shape[0]
   batch_nmse = 0
   for i in range(batch_size):
       batch_nmse += np.linalg.norm(gt[i] - pred[i]) ** 2 / np.linalg.norm(gt[i]) ** 2
   return batch_nmse

def b_psnr(gt, pred):
   """ Compute Peak Signal to Noise Ratio metric (PSNR) """
   batch_size = gt.shape[0]
   batch_psnr = 0
   for i in range(batch_size):
       mse = ((gt[i] - pred[i]) ** 2).mean()
       max_i = gt[i].max()
       s_psnr = 10 * np.log10((max_i ** 2) / mse)
       batch_psnr += s_psnr
   return batch_psnr

def b_ssim(gt, pred):
   """ Compute Structural Similarity Index Metric (SSIM). """
   batch_size = gt.shape[0]
   batch_ssim = 0
   for i in range(batch_size):
       s_ssim = ssim(gt[i], pred[i])
       batch_ssim += s_ssim
   return batch_ssim


def val(model,dataloader,mask_weight):
   """
   Calculate accuracy in val_dataset
   """
   model.eval()
   total_nmse = 0
   total_psnr = 0
   total_ssim = 0
   for ii, (val_input, label, sen, dc_fft) in tqdm(enumerate(dataloader)):
       val_input = val_input.to(opt.device)
       
       label = label.data.numpy() # b*c*W*h*2 real
       val_pred, ks, ke = model(val_input, mask_weight, dc_fft)
       val_pred = val_pred.data.cpu().numpy()
       sen = sen.data.numpy()       

       sen = sen[...,0] + 1j * sen[...,1]       #bcwh imag

       label = label[...,0] + 1j * label[...,1] #bcwh imag
       sen_label = np.absolute(np.sum(label * sen, 1))

       val_pred = val_pred[...,0] + 1j * val_pred[...,1] #bcwh imag
       sen_val_pred = np.absolute(np.sum(val_pred * sen, 1))

       total_nmse += b_nmse(sen_label, sen_val_pred)
       total_psnr += b_psnr(sen_label, sen_val_pred)
       total_ssim += b_ssim(sen_label, sen_val_pred)

   total_nmse /=1410 # 30 slices * 47 patients
   total_psnr /= 1410
   total_ssim /= 1410

   model.train()
   return total_nmse, total_psnr, total_ssim

def help():
   """
   python file.py help
   """
   
   print("""
   usage : python file.py &lt;function&gt; [--args=value]
   &lt;function&gt; := train | test | help
   example: 
           python {0} train --env='env0701' --lr=0.01
           python {0} test --dataset='path/to/dataset/root/'
           python {0} help
   avaiable args:""".format(__file__))

   from inspect import getsource
   source = (getsource(opt.__class__))
   print(source)

if __name__=='__main__':
   # train()
   test()

</code></pre>

ptrblck: <p>Try to remove parts of your code to see if these are responsible for the increase in memory usage. E.g. start with the <code>TensorBoard</code> usage and see if removing it would help.</p>