imaginaire.utils package

Subpackages

• imaginaire.utils.visualization package
  • Submodules
  • imaginaire.utils.visualization.common module
  • imaginaire.utils.visualization.face module
  • imaginaire.utils.visualization.pose module
  • Module contents

Submodules

imaginaire.utils.cudnn module

imaginaire.utils.cudnn.init_cudnn(deterministic, benchmark)

Initialize the cudnn module. The two things to consider is whether to use cudnn benchmark and whether to use cudnn deterministic. If cudnn benchmark is set, then the cudnn deterministic is automatically false.

Parameters

• deterministic (bool) – Whether to use cudnn deterministic.

• benchmark (bool) – Whether to use cudnn benchmark.

imaginaire.utils.data module

Utils for handling datasets.

class imaginaire.utils.data.Augmentor(aug_list, individual_video_frame_aug_list, image_data_types, is_mask, keypoint_data_types, interpolator)

Bases: object

Handles data augmentation using albumentations library.

perform_augmentation(inputs, paired, augment_ops)

Entry point for augmentation.

Parameters

• inputs (dict) – Keys are from self.augmentable_data_types. Values are list of numpy.ndarray (list of images).

• paired (bool) – Apply same augmentation to all input keys?

• augment_ops (list) – The augmentation operations.

imaginaire.utils.data.get_class_number(data_cfg)

Get number of classes for class-conditional GAN model

Parameters

data_cfg (obj) – Data configuration structure.

Returns

Number of classes.

Return type

(int)

imaginaire.utils.data.get_crop_h_w(augmentation)

Get height and width of crop.

Parameters

augmentation (dict) – Dict of applied augmentations.

Returns

• crop_h (int): Height of the image crop.

• crop_w (int): Width of the image crop.

Return type

(dict)

imaginaire.utils.data.get_image_size(x)

imaginaire.utils.data.get_paired_input_image_channel_number(data_cfg)

Get number of channels for the input image.

Parameters

data_cfg (obj) – Data configuration structure.

Returns

Number of input image channels.

Return type

num_channels (int)

imaginaire.utils.data.get_paired_input_label_channel_number(data_cfg, video=False)

Get number of channels for the input label map.

Parameters

• data_cfg (obj) – Data configuration structure.

• video (bool) – Whether we are dealing with video data.

Returns

Number of input label map channels.

Return type

num_channels (int)

imaginaire.utils.data.load_from_folder(keys, handles)

Load keys from lmdb handles.

Parameters

• keys (dict) – This has data_type as key, and a list of paths as values.

• handles (dict) – This has data_type as key, and Folder handle as value.

Returns

This has data_type as key, and a list of decoded items from

folders as value.

Return type

data (dict)

imaginaire.utils.data.load_from_lmdb(keys, lmdbs)

Load keys from lmdb handles.

Parameters

• keys (dict) – This has data_type as key, and a list of paths into LMDB as values.

• lmdbs (dict) – This has data_type as key, and LMDB handle as value.

Returns

This has data_type as key, and a list of decoded items from

LMDBs as value.

Return type

data (dict)

imaginaire.utils.data.load_from_object_store(keys, handles)

Load keys from AWS S3 handles.

Parameters

• keys (dict) – This has data_type as key, and a list of paths as values.

• handles (dict) – This has data_type as key, and Folder handle as value.

Returns

This has data_type as key, and a list of decoded items from

folders as value.

Return type

data (dict)

imaginaire.utils.dataset module

imaginaire.utils.dataset.get_test_dataloader(cfg)

Return dataset objects for testing

Parameters

cfg (obj) – Global configuration file.

Returns

Val data loader. It may not contain the ground truth.

Return type

(obj)

imaginaire.utils.dataset.get_train_and_val_dataloader(cfg, seed=0)

Return dataset objects for the training and validation sets.

Parameters

cfg (obj) – Global configuration file.

Returns

• train_data_loader (obj): Train data loader.

• val_data_loader (obj): Val data loader.

Return type

(dict)

imaginaire.utils.diff_aug module

imaginaire.utils.diff_aug.apply_diff_aug(data, keys, aug_policy, inplace=False, **kwargs)

Applies differentiable augmentation. :param data: Input data. :type data: dict :param keys: Keys to the data values that we want to apply

differentiable augmentation to.

Parameters

aug_policy (str) – Type of augmentation(s), 'color', 'translation', or 'cutout' separated by ','.

imaginaire.utils.diff_aug.diff_aug(x, policy='', channels_first=True, **kwargs)

imaginaire.utils.diff_aug.rand_brightness(x, **kwargs)

imaginaire.utils.diff_aug.rand_contrast(x, **kwargs)

imaginaire.utils.diff_aug.rand_cutout(x, ratio=0.5, **kwargs)

imaginaire.utils.diff_aug.rand_saturation(x, **kwargs)

imaginaire.utils.diff_aug.rand_translation(x, ratio=0.125, **kwargs)

imaginaire.utils.diff_aug.rand_translation_scale(x, trans_r=0.125, scale_r=0.125, mode='bilinear', padding_mode='reflection', **kwargs)

imaginaire.utils.distributed module

imaginaire.utils.distributed.dist_all_gather_tensor(tensor)

gather to all ranks

imaginaire.utils.distributed.dist_all_reduce_tensor(tensor, reduce='mean')

Reduce to all ranks

imaginaire.utils.distributed.dist_reduce_tensor(tensor, rank=0, reduce='mean')

Reduce to rank 0

imaginaire.utils.distributed.get_rank()

Get rank of the thread.

imaginaire.utils.distributed.get_world_size()

Get world size. How many GPUs are available in this job.

imaginaire.utils.distributed.init_dist(local_rank, backend='nccl', **kwargs)

Initialize distributed training

imaginaire.utils.distributed.is_local_master()

imaginaire.utils.distributed.is_master()

check if current process is the master

imaginaire.utils.distributed.master_only(func)

Apply this function only to the master GPU.

imaginaire.utils.distributed.master_only_print(*args)

master-only print

imaginaire.utils.gpu_affinity module

class imaginaire.utils.gpu_affinity.device(device_idx)

Bases: object

Device used for nvml.

getCpuAffinity()

Get CPU affinity

getName()

Get obect name

imaginaire.utils.gpu_affinity.deviceGetCount()

Get number of devices

imaginaire.utils.gpu_affinity.set_affinity(gpu_id=None)

Set GPU affinity

Parameters

gpu_id (int) – Which gpu device.

imaginaire.utils.gpu_affinity.systemGetDriverVersion()

Get Driver Version

imaginaire.utils.init_weight module

imaginaire.utils.init_weight.weights_init(init_type='normal', gain=0.02, bias=None)

Initialize weights in the network.

Parameters

• init_type (str) – The name of the initialization scheme.

• gain (float) – The parameter that is required for the initialization scheme.

• bias (object) – If not None, specifies the initialization parameter for bias.

Returns

init function to be applied.

Return type

(obj)

imaginaire.utils.init_weight.weights_rescale()

imaginaire.utils.io module

imaginaire.utils.io.download_file(URL, destination)

Download a file from google drive or pbss by using the url.

Parameters

• URL – GDrive URL or PBSS pre-signed URL for the checkpoint.

• destination – Path to save the file.

Returns:

imaginaire.utils.io.download_file_from_google_drive(URL, destination)

Download a file from google drive.

Parameters

• URL – GDrive file ID.

• destination – Path to save the file.

Returns:

imaginaire.utils.io.get_checkpoint(checkpoint_path, url='')

Get the checkpoint path. If it does not exist yet, download it from the url.

Parameters

• checkpoint_path (str) – Checkpoint path.

• url (str) – URL to download checkpoint.

Returns

Full checkpoint path.

Return type

(str)

imaginaire.utils.io.get_confirm_token(response)

Get confirm token

Parameters

response – Check if the file exists.

Returns:

imaginaire.utils.io.save_intermediate_training_results(visualization_images, logdir, current_epoch, current_iteration)

Save intermediate training results for debugging purpose.

Parameters

• visualization_images (tensor) – Image where pixel values are in [-1, 1].

• logdir (str) – Where to save the image.

• current_epoch (int) – Current training epoch.

• current_iteration (int) – Current training iteration.

imaginaire.utils.io.save_pilimage_in_jpeg(fullname, output_img)

Save PIL Image to JPEG.

Parameters

• fullname (str) – Full save path.

• output_img (PIL Image) – Image to be saved.

imaginaire.utils.io.save_response_content(response, destination)

Save response content

Parameters

• response –

• destination – Path to save the file.

Returns:

imaginaire.utils.lmdb module

imaginaire.utils.lmdb.build_lmdb(filepaths, keys, output_filepath, map_size, large)

Write out lmdb containing (key, contents of filepath) to file.

Parameters

• filepaths (list) – List of filepath strings.

• keys (list) – List of key strings associated with filepaths.

• output_filepath (str) – Location to write LMDB to.

• map_size (int) – Size of LMDB.

• large (bool) – Is the dataset large?

imaginaire.utils.lmdb.check_and_add(filepath, key, filepaths, keys, remove_missing=False)

Add filepath and key to list of filepaths and keys.

Parameters

• filepath (str) – Filepath to add.

• key (str) – LMDB key for this filepath.

• filepaths (list) – List of filepaths added so far.

• keys (list) – List of keys added so far.

• remove_missing (bool) – If True, removes missing files, otherwise raises an error.

Returns

Size of file at filepath.

Return type

(int)

imaginaire.utils.lmdb.construct_file_path(root, data_type, sequence, filename, ext)

Get file path for our dataset structure.

imaginaire.utils.lmdb.create_metadata(data_root=None, cfg=None, paired=None, input_list='')

Main function.

Parameters

• data_root (str) – Location of dataset root.

• cfg (object) – Loaded config object.

• paired (bool) – Paired or unpaired data.

• input_list (str) – Path to filename containing list of inputs.

Returns

• all_filenames (dict): Key of data type, values with sequences.

• extensions (dict): Extension of each data type.

Return type

(tuple)

imaginaire.utils.lmdb.get_all_filenames_from_list(list_name)

Get all filenames from list.

Parameters

list_name (str) – Path to filename list.

Returns

Folder name for key, and filename for values.

Return type

all_filenames (dict)

imaginaire.utils.lmdb.get_lmdb_data_types(cfg)

Get the data types which should be put in LMDB.

Parameters

cfg – Configuration object.

imaginaire.utils.lmdb.write_entry(txn, key, filepath)

Dump binary contents of file associated with key to LMDB.

Parameters

• txn – handle to LMDB.

• key (str) – LMDB key for this filepath.

• filepath (str) – Filepath to add.

imaginaire.utils.logging module

imaginaire.utils.logging.get_date_uid()

Generate a unique id based on date. :returns: Return uid string, e.g. ‘20171122171307111552’. :rtype: str

imaginaire.utils.logging.init_logging(config_path, logdir)

Create log directory for storing checkpoints and output images.

Parameters

• config_path (str) – Path to the configuration file.

• logdir (str) – Log directory name

Returns

Return log dir

Return type

str

imaginaire.utils.logging.make_logging_dir(logdir)

Create the logging directory

Parameters

logdir (str) – Log directory name

imaginaire.utils.meters module

class imaginaire.utils.meters.Meter(name, reduce=True)

Bases: object

Meter is to keep track of statistics along steps. Meters write values for purpose like printing average values. Meters can be flushed to log files (i.e. TensorBoard for now) regularly.

Parameters

• name (str) – the name of meter

• reduce (bool) – If True, perform a distributed reduce for the log values across all GPUs.

flush(step)

Write the value in the tensorboard.

Parameters

step (int) – Epoch or iteration number.

reset()

Reset the meter values

write(value)

Record the value

write_image(img_grid, step)

Write the value in the tensorboard.

Parameters

• img_grid –

• step (int) – Epoch or iteration number.

imaginaire.utils.meters.get_weight_stats(mod)

Get weight state

Parameters

mod – Pytorch module

imaginaire.utils.meters.set_summary_writer(log_dir)

Set summary writer

Parameters

log_dir (str) – Log directory.

imaginaire.utils.meters.sn_reshape_weight_to_matrix(weight)

Reshape weight to obtain the matrix form.

Parameters

weight (Parameters) – pytorch layer parameter tensor.

imaginaire.utils.meters.write_summary(name, summary, step, hist=False)

Utility function for write summary to log_writer.

imaginaire.utils.misc module

Miscellaneous utils.

imaginaire.utils.misc.apply_imagenet_normalization(input)

Normalize using ImageNet mean and std.

Parameters

input (4D tensor NxCxHxW) – The input images, assuming to be [-1, 1].

Returns

Normalized inputs using the ImageNet normalization.

imaginaire.utils.misc.get_and_setattr(cfg, name, default)

Get attribute with default choice. If attribute does not exist, set it using the default value.

Parameters

• cfg (obj) – Config options.

• name (str) – Attribute name.

• default (obj) – Default attribute.

Returns

Desired attribute.

Return type

(obj)

imaginaire.utils.misc.get_nested_attr(cfg, attr_name, default)

Iteratively try to get the attribute from cfg. If not found, return default.

Parameters

• cfg (obj) – Config file.

• attr_name (str) – Attribute name (e.g. XXX.YYY.ZZZ).

• default (obj) – Default return value for the attribute.

Returns

Attribute value.

Return type

(obj)

imaginaire.utils.misc.gradient_norm(model)

Return the gradient norm of model.

Parameters

model (PyTorch module) – Your network.

imaginaire.utils.misc.random_shift(x, offset=0.05, mode='bilinear', padding_mode='reflection')

Randomly shift the input tensor.

Parameters

• x (4D tensor) – The input batch of images.

• offset (int) – The maximum offset ratio that is between [0, 1].

• maximum shift is offset * image_size for each direction. (The) –

• mode (str) – The resample mode for ‘F.grid_sample’.

• padding_mode (str) – The padding mode for ‘F.grid_sample’.

Returns

The randomly shifted image.

Return type

x (4D tensor)

imaginaire.utils.misc.requires_grad(model, require=True)

Set a model to require gradient or not.

Parameters

• model (nn.Module) – Neural network model.

• require (bool) – Whether the network requires gradient or not.

Returns:

imaginaire.utils.misc.slice_tensor(data, start, end)

Slice all tensors from start to end. :param data: Input data. :type data: dict, list or tensor

imaginaire.utils.misc.split_labels(labels, label_lengths)

Split concatenated labels into their parts.

Parameters

• labels (torch.Tensor) – Labels obtained through concatenation.

• label_lengths (OrderedDict) – Containing order of labels & their lengths.

Returns:

imaginaire.utils.misc.to_channels_last(data)

Move all data to channels_last format.

Parameters

data (dict, list or tensor) – Input data.

imaginaire.utils.misc.to_cpu(data)

Move all tensors inside data to cpu.

Parameters

data (dict, list, or tensor) – Input data.

imaginaire.utils.misc.to_cuda(data)

Move all tensors inside data to gpu.

Parameters

data (dict, list, or tensor) – Input data.

imaginaire.utils.misc.to_device(data, device)

Move all tensors inside data to device.

Parameters

• data (dict, list, or tensor) – Input data.

• device (str) – ‘cpu’ or ‘cuda’.

imaginaire.utils.misc.to_float(data)

Move all halfs to float.

Parameters

data (dict, list or tensor) – Input data.

imaginaire.utils.misc.to_half(data)

Move all floats to half.

Parameters

data (dict, list or tensor) – Input data.

imaginaire.utils.model_average module

class imaginaire.utils.model_average.ModelAverage(module, beta=0.9999, start_iteration=1000, remove_wn_wrapper=True)

Bases: torch.nn.modules.module.Module

In this model average implementation, the spectral layers are absorbed in the model parameter by default. If such options are turned on, be careful with how you do the training. Remember to re-estimate the batch norm parameters before using the model.

Parameters

• module (torch nn module) – Torch network.

• beta (float) – Moving average weights. How much we weight the past.

• start_iteration (int) – From which iteration, we start the update.

• remove_sn (bool) – Whether we remove the spectral norm when we it.

copy_s2t()

Copy state_dictionary from source to target. Here source is the regular module and the target is the moving average module. Basically, we will copy weights in the regular module to the moving average module.

copy_t2s()

Copy the original weights to the moving average weights.

forward(*inputs, **kwargs)

PyTorch module forward function overload.

sn_compute_weight(weight, u, v)

Compute the spectral norm normalized matrix.

Parameters

• weight (Parameters) – pytorch layer parameter tensor.

• u (tensor) – left singular vectors.

• v (tensor) –

Returns

weight parameter object.

Return type

(Parameters)

sn_reshape_weight_to_matrix(weight)

Reshape weight to obtain the matrix form.

Parameters

weight (Parameters) – pytorch layer parameter tensor.

Returns

Reshaped weight matrix

Return type

(Parameters)

training = None

update_average()

Update the moving average.

imaginaire.utils.model_average.calibrate_batch_norm_momentum(m)

Calibrate batch norm momentum

Parameters

m – Pytorch module

imaginaire.utils.model_average.reset_batch_norm(m)

Reset batch norm statistics

Parameters

m – Pytorch module

imaginaire.utils.path module

Utils to deal with directories and paths.

imaginaire.utils.path.get_immediate_subdirectories(input_dir)

List dirs immediately under input_dir.

Parameters

• input_dir (str) – Directory to list children of.

• Returns –

• (list) – List of directory paths relative to input_dir.

imaginaire.utils.path.get_recursive_subdirectories(input_dir, ext)

List dirs recursively under input_dir.

Parameters

• input_dir (str) – Directory to list children of.

• ext (str) – Extension of files expected in this directory.

• Returns –

• (list) – List of directory paths relative to input_dir.

imaginaire.utils.trainer module

class imaginaire.utils.trainer.WrappedModel(module)

Bases: torch.nn.modules.module.Module

Dummy wrapping the module.

forward(*args, **kwargs)

PyTorch module forward function overload.

training = None

imaginaire.utils.trainer.get_model_optimizer_and_scheduler(cfg, seed=0)

Return the networks, the optimizers, and the schedulers. We will first set the random seed to a fixed value so that each GPU copy will be initialized to have the same network weights. We will then use different random seeds for different GPUs. After this we will wrap the generator with a moving average model if applicable. It is followed by getting the optimizers and data distributed data parallel wrapping.

Parameters

• cfg (obj) – Global configuration.

• seed (int) – Random seed.

Returns

• net_G (obj): Generator network object.

• net_D (obj): Discriminator network object.

• opt_G (obj): Generator optimizer object.

• opt_D (obj): Discriminator optimizer object.

• sch_G (obj): Generator optimizer scheduler object.

• sch_D (obj): Discriminator optimizer scheduler object.

Return type

(dict)

imaginaire.utils.trainer.get_optimizer(cfg_opt, net)

Return the scheduler object.

Parameters

• cfg_opt (obj) – Config for the specific optimization module (gen/dis).

• net (obj) – PyTorch network object.

Returns

Pytorch optimizer

Return type

(obj)

imaginaire.utils.trainer.get_optimizer_for_params(cfg_opt, params)

Return the scheduler object.

Parameters

• cfg_opt (obj) – Config for the specific optimization module (gen/dis).

• params (obj) – Parameters to be trained by the parameters.

Returns

Optimizer

Return type

(obj)

imaginaire.utils.trainer.get_scheduler(cfg_opt, opt)

Return the scheduler object.

Parameters

• cfg_opt (obj) – Config for the specific optimization module (gen/dis).

• opt (obj) – PyTorch optimizer object.

Returns

Scheduler

Return type

(obj)

imaginaire.utils.trainer.get_trainer(cfg, net_G, net_D=None, opt_G=None, opt_D=None, sch_G=None, sch_D=None, train_data_loader=None, val_data_loader=None)

Return the trainer object.

Parameters

• cfg (Config) – Loaded config object.

• net_G (obj) – Generator network object.

• net_D (obj) – Discriminator network object.

• opt_G (obj) – Generator optimizer object.

• opt_D (obj) – Discriminator optimizer object.

• sch_G (obj) – Generator optimizer scheduler object.

• sch_D (obj) – Discriminator optimizer scheduler object.

• train_data_loader (obj) – Train data loader.

• val_data_loader (obj) – Validation data loader.

Returns

Trainer object.

Return type

(obj)

imaginaire.utils.trainer.set_random_seed(seed, by_rank=False)

Set random seeds for everything.

Parameters

• seed (int) – Random seed.

• by_rank (bool) –

imaginaire.utils.trainer.wrap_model_and_optimizer(cfg, net_G, net_D, opt_G, opt_D)

Wrap the networks and the optimizers with AMP DDP and (optionally) model average.

Parameters

• cfg (obj) – Global configuration.

• net_G (obj) – Generator network object.

• net_D (obj) – Discriminator network object.

• opt_G (obj) – Generator optimizer object.

• opt_D (obj) – Discriminator optimizer object.

Returns

• net_G (obj): Generator network object.

• net_D (obj): Discriminator network object.

• opt_G (obj): Generator optimizer object.

• opt_D (obj): Discriminator optimizer object.

Return type

(dict)

Module contents