Implementation:Junyanz Pytorch CycleGAN and pix2pix Visualizer

Field	Value
source	pytorch-CycleGAN-and-pix2pix\|https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix
domains	Training, Monitoring

Overview

Concrete tool for logging training images and losses during GAN training provided by the pytorch-CycleGAN-and-pix2pix framework.

The Visualizer class (defined in util/visualizer.py, lines 40--168) is the single entry point for all training-time monitoring. It unifies four output channels -- console printing, file logging, HTML gallery generation, and Weights & Biases cloud logging -- behind a small method API that the training loop calls at configurable intervals. A companion standalone function save_images (lines 12--37) handles saving inference results to HTML during test time.

Code Reference

Source file: util/visualizer.py

Import

from util.visualizer import Visualizer

Class Signature

class Visualizer:
    """Display/save images and print/save logging information.

    Uses wandb for logging (optional) and Python library 'dominate'
    (wrapped in 'HTML') for creating HTML files with images.
    """

    def __init__(self, opt):
        """Initialize the Visualizer class.

        Parameters:
            opt -- stores all experiment flags; subclass of BaseOptions
        """

Source location: Lines 40--168

Constructor Detail (`init`, L46--83)

The constructor performs four sequential steps:

Cache options -- Stores opt, extracts opt.isTrain, opt.no_html, opt.display_winsize, opt.name, opt.use_wandb.
Initialize WandB -- If opt.use_wandb is True and the process is rank 0 (or DDP is not initialized), calls wandb.init(project=opt.wandb_project_name, name=opt.name, config=opt). On non-zero ranks, sets self.wandb_run = None.
Create HTML directory -- If training with HTML enabled, creates <checkpoints_dir>/<name>/web/ and .../web/images/.
Create loss log file -- Opens <checkpoints_dir>/<name>/loss_log.txt in append mode and writes a timestamped header line.

Method: `reset` (L85--87)

def reset(self):
    """Reset the self.saved status"""
    self.saved = False

Resets the self.saved flag so that the next call to display_current_results will save images to the HTML gallery. Called once at the start of each epoch.

Method: `display_current_results` (L98--132)

def display_current_results(self, visuals, epoch: int, total_iters: int, save_result=False):
    """Save current results to wandb and HTML file."""

Parameters:

Parameter	Type	Description
`visuals`	`OrderedDict[str, Tensor]`	Named image tensors (e.g., `real_A`, `fake_B`, `rec_A`)
`epoch`	`int`	Current training epoch number
`total_iters`	`int`	Cumulative iteration count across all epochs
`save_result`	`bool`	If `True`, force save to HTML even if already saved this epoch

Behavior:

Skips entirely on non-zero ranks in distributed training.
WandB path: Converts each visual tensor to a NumPy image via util.tensor2im, wraps it in wandb.Image with a caption, and logs the dictionary under the results/ namespace at the current total_iters step.
HTML path: If HTML is enabled and images have not yet been saved this epoch (or save_result=True):
- Saves each visual as epoch{NNN}_{label}.png to the images directory.
- Rebuilds the entire HTML page from epoch N down to epoch 1, with each epoch as a header row and its images in a table.
- Calls webpage.save() to write index.html.

Method: `plot_current_losses` (L134--146)

def plot_current_losses(self, total_iters, losses):
    """Log current losses to wandb.

    Parameters:
        total_iters (int)     -- current training iteration
        losses (OrderedDict)  -- (name, float) pairs
    """

On rank 0 with WandB enabled, calls self.wandb_run.log(losses, step=total_iters). This logs scalar loss values (e.g., G_GAN, G_L1, D_real, D_fake) directly as WandB step metrics.

Method: `print_current_losses` (L148--168)

def print_current_losses(self, epoch, iters, losses, t_comp, t_data):
    """Print current losses on console; also save to disk.

    Parameters:
        epoch (int)           -- current epoch
        iters (int)           -- current iteration within this epoch
        losses (OrderedDict)  -- (name, float) pairs
        t_comp (float)        -- computation time per data point
        t_data (float)        -- data loading time per data point
    """

Behavior:

Reads LOCAL_RANK from the environment (defaults to 0).
Formats a message string: [Rank N] (epoch: E, iters: I, time: T, data: D), loss1: V, loss2: V, ...
Prints the message to console on all ranks (so each GPU's losses are visible).
On rank 0 only, appends the message to loss_log.txt.

Standalone Function: `save_images` (L12--37)

def save_images(webpage, visuals, image_path, aspect_ratio=1.0, width=256):
    """Save images to the disk.

    Parameters:
        webpage (HTML)           -- the HTML webpage instance
        visuals (OrderedDict)    -- (name, tensor/numpy) pairs
        image_path (str)         -- used to derive the image file name
        aspect_ratio (float)     -- aspect ratio of saved images
        width (int)              -- display width in the HTML table
    """

This function is used during test/inference (not training). It:

Extracts the stem of image_path[0] as the base name.
Adds a header row to the HTML page with that name.
For each visual, converts the tensor to a NumPy image, saves it as {name}_{label}.png, and adds it to the HTML table row.

I/O Contract

Inputs

Input	Type	Description
`opt`	`argparse.Namespace`	Full experiment options. Required fields: `isTrain`, `no_html`, `display_winsize`, `name`, `checkpoints_dir`, `use_wandb`, `wandb_project_name`.
`visuals`	`OrderedDict[str, Tensor]`	Named image tensors produced by the model (e.g., `real_A`, `fake_B`, `rec_A`, `idt_A`).
`losses`	`OrderedDict[str, float]`	Named scalar losses (e.g., `G_GAN`, `G_L1`, `D_real`, `D_fake`).
`epoch`	`int`	Current epoch number (1-indexed).
`iters`	`int`	Current iteration within the epoch.
`total_iters`	`int`	Cumulative iteration count.
`t_comp`	`float`	Computation time per data point (seconds).
`t_data`	`float`	Data loading time per data point (seconds).

Outputs

Output	Location	Description
Console logs	`stdout`	Formatted loss strings with rank, epoch, iteration, and timing info. Printed on all ranks.
`loss_log.txt`	`<checkpoints_dir>/<name>/loss_log.txt`	Persistent append-only log of all loss messages. Written on rank 0 only.
HTML gallery	`<checkpoints_dir>/<name>/web/index.html`	Auto-refreshing HTML page with epoch-by-epoch image tables. Images stored in `.../web/images/`.
WandB logs	WandB cloud project	Loss scalars logged as step metrics; generated images logged as `wandb.Image` under the `results/` namespace. Logged on rank 0 only.

Usage Examples

Example 1: Standard Training Loop Integration

from options.train_options import TrainOptions
from data import create_dataset
from models import create_model
from util.visualizer import Visualizer

opt = TrainOptions().parse()
dataset = create_dataset(opt)
model = create_model(opt)
model.setup(opt)

visualizer = Visualizer(opt)
total_iters = 0

for epoch in range(opt.epoch_count, opt.n_epochs + opt.n_epochs_decay + 1):
    visualizer.reset()

    for i, data in enumerate(dataset):
        total_iters += opt.batch_size
        epoch_iter = i * opt.batch_size

        model.set_input(data)
        model.optimize_parameters()

        # Display generated images at display_freq intervals
        if total_iters % opt.display_freq == 0:
            save_result = total_iters % opt.update_html_freq == 0
            model.compute_visuals()
            visualizer.display_current_results(
                model.get_current_visuals(), epoch, total_iters, save_result
            )

        # Print losses to console and log to WandB at print_freq intervals
        if total_iters % opt.print_freq == 0:
            losses = model.get_current_losses()
            t_comp = 0.05  # seconds per data point
            t_data = 0.01  # seconds per data point
            visualizer.print_current_losses(epoch, epoch_iter, losses, t_comp, t_data)
            visualizer.plot_current_losses(total_iters, losses)

Example 2: Using save_images for Test/Inference

from util.visualizer import save_images
from util import html as html_util

webpage = html_util.HTML(web_dir="results/my_test/web", title="Test Results")

# Inside the test loop:
for i, data in enumerate(dataset):
    model.set_input(data)
    model.test()
    visuals = model.get_current_visuals()
    img_path = model.get_image_paths()
    save_images(webpage, visuals, img_path, aspect_ratio=1.0, width=256)

webpage.save()

Example 3: Enabling WandB Logging

Launch training with the --use_wandb flag:

python train.py \
    --dataroot ./datasets/maps \
    --name maps_cyclegan \
    --model cycle_gan \
    --use_wandb \
    --wandb_project_name "my-cyclegan-experiments"

This initializes a WandB run on rank 0. Loss scalars and generated images are automatically logged at the configured intervals. Visit the WandB project dashboard to compare runs, overlay loss curves, and browse image panels.

Related Pages

Page Connections

Double-click a node to navigate. Hold to expand connections.

Principle

Implementation

Heuristic

Environment