Out of CPU Memory caused by Parallel()

[Wenwen717]

First of all, thank you again for open-sourcing such excellent work.
I am trying to run blending_train.py, but when I reach the following section:

class Blending_dataset(Dataset):
    def __init__(self, exps, path, net_trainer):
        super().__init__()
        downsample_256 = BicubicDownSample(factor=4)
        data = Parallel(n_jobs=1)(
            delayed(prepare_item)(exp, path) for (p1, p2, p3) in tqdm(exps) for exp in [(p1, p2, p3), (p1, p3, p2)])**
        data = [elem for elem in data if elem is not None]
        print(f'Load: {len(data)}/{2 * len(exps)}', file=sys.stderr)

I notice that my CPU memory gradually gets filled. Can I solve this issue by configuring a parameter in the Parallel class?

class Parallel(Logger):
    def __init__(self, n_jobs=None, backend=None, verbose=0, timeout=None,
                 pre_dispatch='2 * n_jobs', batch_size='auto',
                 temp_folder=None, max_nbytes='1M', mmap_mode='r',
                 prefer=None, require=None):

I tried reducing n_jobs, batch_size, and max_nbytes, but it doesn't seem to work.

Heartfelt thanks for any suggestions.

[Wenwen717]

I solved this problem by storing the processing results into specified files in batches and then reading them sequentially.

The code is as follows:

def process_in_batches(exps, path, batch_size, temp_dir):
    # Check if the processed temporary files already exist
    existing_files = sorted([os.path.join(temp_dir, f) for f in os.listdir(temp_dir) if f.endswith('.pkl')])
    if existing_files:
        print(f"Using existing batch files from {temp_dir}")
        return existing_files

    if not os.path.exists(temp_dir):
        os.makedirs(temp_dir)
    temp_files = []
    for i in tqdm(range(0, len(exps), batch_size)):
        batch = exps[i:i + batch_size]
        batch_results = Parallel(n_jobs=-1)(
            delayed(prepare_item)(exp, path) for (p1, p2, p3) in batch for exp in [(p1, p2, p3), (p1, p3, p2)])
        temp_file = os.path.join(temp_dir, f'batch_{i}.pkl')
        dump(batch_results, temp_file)
        temp_files.append(temp_file)

        # Clean up intermediate results
        del batch_results
        gc.collect()

    return temp_files

def load_results_generator(temp_files):
    for temp_file in temp_files:
        for result in load(temp_file):
            yield result

class Blending_dataset(Dataset):
    def __init__(self, exps, path, net_trainer):
        super().__init__()
        downsample_256 = BicubicDownSample(factor=4)

        batch_size = 100
        temp_dir = 'temp_results' 
        temp_files = process_in_batches(exps, path, batch_size, temp_dir)

        self.items = []
        with torch.no_grad():
            for temp_file in tqdm(temp_files):
                data = [elem for elem in load(temp_file) if elem is not None]
                tmp_dataloader = DataLoader(data, batch_size=24, pin_memory=False, shuffle=False)
                for (Color_S, Align_S, Align_F, Color_I, Face_I) in tmp_dataloader:
                    HM_3D, HM_3E = net_trainer.generate_mask(Color_I.to('cuda'))

                    HM_1D, _ = net_trainer.generate_mask(Face_I.to('cuda'))
                    I_X, _ = net_trainer.net.generator([Align_S.to('cuda')], input_is_latent=True, return_latents=False,
                                                       start_layer=4,
                                                       end_layer=8, layer_in=Align_F.to('cuda'))
                    HM_XD, HM_XE = net_trainer.generate_mask(I_X)

                    target_mask = ((1 - HM_1D) * (1 - HM_3D) * (1 - HM_XD)).cpu()
                    HM_3E = HM_3E.cpu()
                    HM_XE = HM_XE.cpu()
                    self.items.extend(
                        [item for item in zip(*list(map(lambda x: [item.squeeze(0) for item in torch.split(x, 1)],
                                                        (Color_S,
                                                         Align_S,
                                                         Align_F,
                                                         downsample_256(Color_I.to('cuda')).cpu(),
                                                         downsample_256(Face_I.to('cuda')).cpu(),
                                                         target_mask, HM_3E, HM_XE)))
                                              ) if item[-2].any() and item[-1].any()]
                    )

                # Clean up intermediate results
                del data
                gc.collect()

        print(f'dataset: {len(self.items)}/{len(temp_files) * batch_size}', file=sys.stderr)

But I'm not sure if you've used a more efficient method to solve the problem, such as by changing some parameters.

Thank you for any response.

[maximkm]

Hello,
to maximize the utilization of the training, we used datasets that are fully loaded into RAM. On our cluster, no problems were encountered.

In case you are facing the problem of running out of RAM even for items, probably the most efficient solution is to save all the items to disk and in __getitem__ read it.

[maximkm]

In a future version of the scripts, we will significantly rework the training to make such problems easier to fix.