this is could reproduce , gpu memory continue increase
import os
import sys
import json
import torch
import random
import resource
from torch.utils.data import IterableDataset
from torch.utils.data import DataLoader
from itertools import tee

UNK_ID=100
BOS_ID=101

class CustomIterableDataset(IterableDataset):
    def __init__(self, task_def, task_id, batch_size=32,
                 gpu=True, is_train=True, epochs=10,
                 maxlen=128, dropout_w=0.005):
        super(CustomIterableDataset).__init__()
        self.task_def = task_def
        self.task_id = task_id
        self.batch_size = batch_size
        self.maxlen = maxlen
        self.is_train = is_train
        self.epochs = 1 if not is_train else epochs
        self.gpu = gpu
        self.dropout_w = dropout_w
        self.pairwise_size = 1


    def patch(self, v):
        v = v.cuda(non_blocking=True)
        return v

    def _dir_iter(self):
        for i in range(0,3):
            lines = []
            reader = [i for i in range(0,1000000)]
            for line in reader:
                line = torch.LongTensor(self.batch_size, 10).fill_(0)
                if len(lines) >= self.batch_size:
                    yield lines
                    lines = []
                lines.append(self.patch(line))
            del reader
            yield lines

    def __iter__(self):
        mapped_itr =self._dir_iter()
        return mapped_itr


if __name__ == '__main__':
    dataset_iter = CustomIterableDataset(task_def=None, task_id=None)
    data_generator = DataLoader(dataset_iter, batch_size=None)
    train_generator_list = []
    train_generator_list.append(iter(data_generator))
    copy_iter_list = []
    for id, first_it in enumerate(train_generator_list):
        first_itr, second_itr = tee(first_it)
        train_generator_list[id] = first_itr
        copy_iter_list.append(second_itr)
    i= 0
    while True:
        try:
            i += 1
            batch_data = next(train_generator_list[0])
            if i % 100 == 0:
                max_mem_used = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
                print("Mem: {:.2f} MB".format(max_mem_used / 1024))
        except StopIteration:
            # end of one epoch
            print('again')
            _, train_generator_list[0] = tee(copy_iter_list[0])


image2764×462 137 KB