EnlightenGAN Documentation

To set up EnlightenGAN, ensure you are using python3.5. The project requires significant hardware resources; it is recommended to use at least three NVIDIA 1080ti GPUs, or you must manually adjust the batch size in the configuration to fit your hardware.

Follow these steps to install dependencies and prepare the model directory:

1. Install required Python packages using pip.
2. Create a model directory.
3. Download the VGG pretrained model and place it inside the model directory.

pip install -r requirement.txt
mkdir model

To run inference (testing) on your own images, follow these steps:

1. Download the pretrained model and place it in the ./checkpoints/enlightening directory.
2. Prepare test datasets: Create the following directory structure:
   • ../test_dataset/testA (Place your target test images here)
   • ../test_dataset/testB (Place at least one dummy image here to ensure the program initializes correctly)
3. Run prediction: Execute the script with the --predict flag.

python scripts/script.py --predict

EnlightenGAN uses unpaired images for training and specific datasets for testing.

• Training Data: Unpaired images collected from multiple datasets.
• Testing Data: Includes LIME, MEF, NPE, VV, and DICP datasets.

Datasets can be found via the provided Google Drive links or BaiduYun (as mentioned in the repository issues).

To train the model, you must first launch a visdom.server to enable visualization. Use the following command to start the server in the background:

nohup python -m visdom.server -port=8097

Once the server is running, initiate the training process using the provided script with the --train flag:

python scripts/script.py --train

Setting pin_memory=True in the DataLoader instructs the loader to copy tensors into CUDA pinned memory before returning them. This can significantly speed up the transfer of data from CPU to GPU. This feature is only active if torch.cuda.is_available() is true.

dataloader = DataLoader(dataset, batch_size=32, pin_memory=True)

When using num_workers > 0, each worker is seeded with base_seed + worker_id. You can provide a worker_init_fn to perform additional initialization (like setting NumPy or other library seeds) inside each worker subprocess.

Warning: If using the 'spawn' start method, worker_init_fn cannot be an unpicklable object like a lambda function.

import numpy as np
import torch

def my_init_fn(worker_id):
    # Set seed for numpy in the worker process
    worker_seed = torch.initial_seed() % 2**32
    np.random.seed(worker_seed)

dataloader = DataLoader(
    dataset=my_dataset,
    num_workers=4,
    worker_init_fn=my_init_fn
)

The ConcatDataset class allows you to combine multiple Dataset instances into a single dataset. The concatenation is performed on-the-fly, making it efficient for large-scale datasets. You can also use the + operator between two Dataset objects to achieve the same result via the __add__ method.

from lib.utils.data.dataset import ConcatDataset, TensorDataset
import torch

dataset1 = TensorDataset(torch.randn(5, 3), torch.randn(5))
dataset2 = TensorDataset(torch.randn(10, 3), torch.randn(10))

# Method 1: Explicit instantiation
combined_dataset = ConcatDataset([dataset1, dataset2])

# Method 2: Using the + operator
combined_dataset_alt = dataset1 + dataset2

print(len(combined_dataset))  # Output: 15

The default_collate function merges a list of samples into a mini-batch. It automatically converts various data types into PyTorch tensors:

• Tensors: Stacked into a single tensor. If in a background process, it uses shared memory to avoid extra copies.
• NumPy arrays: Converted to tensors (throws TypeError for string/object arrays).
• Integers/Floats: Converted to LongTensor or DoubleTensor.
• Mappings (dicts): Recursively collates values for each key.
• Sequences (lists/tuples): Recursively collates elements.
• Strings: Returned as a list of strings.

from lib.utils.data.dataloader import default_collate

# Example of what default_collate does internally
batch = [torch.tensor([1, 2]), torch.tensor([3, 4])]
collated = default_collate(batch)
# collated is now torch.tensor([[1, 2], [3, 4]])

The user_scattered_collate(batch) function is a placeholder/identity function intended to be used as a collate_fn in a PyTorch DataLoader. It is designed to work in conjunction with UserScatteredDataParallel to ensure that the data structure passed to the model is compatible with the custom scattering mechanism.

from torch.utils.data import DataLoader
from lib.nn.parallel.data_parallel import user_scattered_collate

# Use it in your DataLoader to prepare batches for UserScatteredDataParallel
dataloader = DataLoader(dataset, batch_size=32, collate_fn=user_scattered_collate)