Running out of CUDA memory

c.justin: <p>Hi there,</p>
<p>I’m new to pytorch and trying to train a GNN with my limited knowledge…However, I met a problem with CUDA running out of memory. My dataset contains graphs with variable sizes.<br>
By printing out the memory allocated by Cuda and the memory reserved by Cuda, it seems if I don’t use a sampler in DataLoader, the memory reserved stays unchanged, but if I use a sampler (with replacement = True), the memory reserved is different, and keep increasing. I’m guessing that if I unluckily sampled a bunch of large graphs needs a lot more memory. However, why after each iteration, those memory are not freed? Is this because I used the sampler in a wrong way?<br>
I’m currently using it by:<br>
loader = GraphDataLoader(ds_train, batch_size=batch_size, drop_last=True, sampler=sampler)</p>
<p>Here is the two memory usage comparison:<br>
Without sampler:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/a/1/a1fe365ce916d9044e348a354346dd6f63a29bb4.png" data-download-href="https://discuss.pytorch.org/uploads/default/a1fe365ce916d9044e348a354346dd6f63a29bb4" title="Screen Shot 2020-04-10 at 2.06.32 PM"><img src="https://discuss.pytorch.org/uploads/default/original/3X/a/1/a1fe365ce916d9044e348a354346dd6f63a29bb4.png" alt="Screen Shot 2020-04-10 at 2.06.32 PM" data-base62-sha1="n73wfEz7unb9cxCvQ12mHIVD3RW" width="245" height="500" data-dominant-color="0B1807"><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">Screen Shot 2020-04-10 at 2.06.32 PM</span><span class="informations">356×724 15.9 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><br>
With sampler:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/3/6/36b1ceaf8f3c2949b22e03594cca33ff2426e967.png" data-download-href="https://discuss.pytorch.org/uploads/default/36b1ceaf8f3c2949b22e03594cca33ff2426e967" title="Screen Shot 2020-04-10 at 2.04.34 PM"><img src="https://discuss.pytorch.org/uploads/default/original/3X/3/6/36b1ceaf8f3c2949b22e03594cca33ff2426e967.png" alt="Screen Shot 2020-04-10 at 2.04.34 PM" data-base62-sha1="7NQJvdZuVdsK96B2bv6KP7OmO5p" width="196" height="500" data-dominant-color="0A1506"><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">Screen Shot 2020-04-10 at 2.04.34 PM</span><span class="informations">410×1042 21.6 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>Thanks for your time looking at this!</p>

ptrblck: <p>Could you post the sampler and your implementation of <code>GraphDataLoader</code>, so that we can have a look, please?</p>

c.justin: <p>Hi there,<br>
Thanks for replying!<br>
GraphDataLoader is just the DataLoader from torch_geometric and Sampler is the torch.utils.data’ s WeightedRandomSampler</p>
<p>Here is the code:</p>
<pre><code class="lang-auto">from torch_geometric.data import DataLoader as GraphDataLoader
from torch_geometric.data import Data
from automate.nn.mate_dataset import RawMateDataset
import automate
from torch.utils.data import Dataset, WeightedRandomSampler
import torch
import numpy as np
from automate.nn.modules.graph_dataset import GraphDataset
from automate.nn.modules.graph_classifier import GraphClassifier


def train_GNN(
    path,
    filters,
    experiment_name,
    save_dir,
    batch_size = 32,
    n_examples = -1,
    lr = 0.001,
    save = False):
    print('creating raw mate dataset')
    ds_raw = automate.nn.mate_dataset.RawMateDataset({'mate_list':path, 'filters':filters})
    print('creating graph dataset')
    ds_graph = GraphDataset(ds_raw)
    if n_examples == -1:
        n_examples = len(ds_graph)
    n_examples = int(.95*n_examples)
    ds_train, ds_val = torch.utils.data.random_split(ds_graph, (n_examples, len(ds_graph) - n_examples))
    print(f'have {len(ds_train)} training examples')
    print(f'have {len(ds_val)} training examples')
    train_weights = ds_raw.weights[ds_train.indices]
    sampler = WeightedRandomSampler(train_weights, len(train_weights))
    print('created loader')
    loader = GraphDataLoader(ds_train, batch_size=batch_size, drop_last=True, sampler=sampler)
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    classifier = GraphClassifier(v_in=71, e_in=6, v_g=30, e_g=6, v_out=200, mc_out=22, i_types=11).float().to(device)
    optimizer = torch.optim.Adam(classifier.parameters(), lr=lr)
    loss = torch.nn.CrossEntropyLoss()
    
    best = -1
    
    print('starting training')
    classifier.train()
    for epoch in range(1000):
        
        try:
            if save and epoch % 100 == 0 and epoch != 0:
                save_model(classifier, save_dir, epoch)
            epoch_loss = 0
            total = 0
            correct = 0
            classifier.train()
            for i, data in enumerate(loader):
                optimizer.zero_grad()
                data = data.to(device)
                target = torch.tensor(data.mate_type).to(device)
                preds = classifier(data)
                err = loss(preds, target)
                _, preds_temp = torch.max(preds.data, 1)
                total += len(target)
                correct += (preds_temp == target).sum().item()
                epoch_loss += err.item()
                err.backward()
                optimizer.step()
            print(f'Epoch {epoch + 1} Loss = {epoch_loss/(i+1)} Train Accuracy = {correct / total}') 
</code></pre>