# ClimateGAN - [ClimateGAN](#climategan) - [Setup](#setup) - [Coding conventions](#coding-conventions) - [updates](#updates) - [interfaces](#interfaces) - [Logging on comet](#logging-on-comet) - [Resources](#resources) - [Example](#example) - [Release process](#release-process) ## Setup **`PyTorch >= 1.1.0`** otherwise optimizer.step() and scheduler.step() are in the wrong order ([docs](https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate)) **pytorch==1.6** to use pytorch-xla or automatic mixed precision (`amp` branch). Configuration files use the **YAML** syntax. If you don't know what `&` and `<<` mean, you'll have a hard time reading the files. Have a look at: * https://dev.to/paulasantamaria/introduction-to-yaml-125f * https://stackoverflow.com/questions/41063361/what-is-the-double-left-arrow-syntax-in-yaml-called-and-wheres-it-specced/41065222 **pip** ``` $ pip install comet_ml scipy opencv-python torch torchvision omegaconf==1.4.1 hydra-core==0.11.3 scikit-image imageio addict tqdm torch_optimizer ``` ## Coding conventions * Tasks * `x` is an input image, in [-1, 1] * `s` is a segmentation target with `long` classes * `d` is a depth map target in R, may be actually `log(depth)` or `1/depth` * `m` is a binary mask with 1s where water is/should be * Domains * `r` is the *real* domain for the masker. Input images are real pictures of urban/suburban/rural areas * `s` is the *simulated* domain for the masker. Input images are taken from our Unity world * `rf` is the *real flooded* domain for the painter. Training images are pairs `(x, m)` of flooded scenes for which the water should be reconstructed, in the validation data input images are not flooded and we provide a manually labeled mask `m` * `kitti` is a special `s` domain to pre-train the masker on [Virtual Kitti 2](https://europe.naverlabs.com/research/computer-vision/proxy-virtual-worlds-vkitti-2/) * it alters the `trainer.loaders` dict to select relevant data sources from `trainer.all_loaders` in `trainer.switch_data()`. The rest of the code is identical. * Flow * This describes the call stack for the trainers standard training procedure * `train()` * `run_epoch()` * `update_G()` * `zero_grad(G)` * `get_G_loss()` * `get_masker_loss()` * `masker_m_loss()` -> masking loss * `masker_s_loss()` -> segmentation loss * `masker_d_loss()` -> depth estimation loss * `get_painter_loss()` -> painter's loss * `g_loss.backward()` * `g_opt_step()` * `update_D()` * `zero_grad(D)` * `get_D_loss()` * painter's disc losses * `masker_m_loss()` -> masking AdvEnt disc loss * `masker_s_loss()` -> segmentation AdvEnt disc loss * `d_loss.backward()` * `d_opt_step()` * `update_learning_rates()` -> update learning rates according to schedules defined in `opts.gen.opt` and `opts.dis.opt` * `run_validation()` * compute val losses * `eval_images()` -> compute metrics * `log_comet_images()` -> compute and upload inferences * `save()` ### Resuming Set `train.resume` to `True` in `opts.yaml` and specify where to load the weights: Use a config's `load_path` namespace. It should have sub-keys `m`, `p` and `pm`: ```yaml load_paths: p: none # Painter weights m: none # Masker weights pm: none # Painter + Masker weights (single ckpt for both) ``` 1. any path which leads to a dir will be loaded as `path / checkpoints / latest_ckpt.pth` 2. if you want to specify a specific checkpoint (not the latest), it MUST be a `.pth` file 3. resuming a `P` **OR** an `M` model, you may only specify 1 of `load_path.p` **OR** `load_path.m`. You may also leave **BOTH** at `none`, in which case `output_path / checkpoints / latest_ckpt.pth` will be used 4. resuming a P+M model, you may specify (`p` AND `m`) **OR** `pm` **OR** leave all at `none`, in which case `output_path / checkpoints / latest_ckpt.pth` will be used to load from a single checkpoint ### Generator * **Encoder**: `trainer.G.encoder` Deeplabv2 or v3-based encoder * Code borrowed from * https://github.com/valeoai/ADVENT/blob/master/advent/model/deeplabv2.py * https://github.com/CoinCheung/DeepLab-v3-plus-cityscapes * **Decoders**: * `trainer.G.decoders["s"]` -> *Segmentation* -> DLV3+ architecture (ASPP + Decoder) * `trainer.G.decoders["d"]` -> *Depth* -> ResBlocks + (Upsample + Conv) * `trainer.G.decoders["m"]` -> *Mask* -> ResBlocks + (Upsample + Conv) -> Binary mask: 1 = water should be there * `trainer.G.mask()` predicts a mask and optionally applies `sigmoid` from an `x` input or a `z` input * **Painter**: `trainer.G.painter` -> [GauGAN SPADE-based](https://github.com/NVlabs/SPADE) * input = masked image * `trainer.G.paint(m, x)` higher level function which takes care of masking * If `opts.gen.p.paste_original_content` the painter should only create water and not reconstruct outside the mask: the output of `paint()` is `painted * m + x * (1 - m)` High level methods of interest: * `trainer.infer_all()` creates a dictionary of events with keys `flood` `wildfire` and `smog`. Can take in a single image or a batch, of numpy arrays or torch tensors, on CPU/GPU/TPU. This method calls, amongst others: * `trainer.G.encode()` to compute the shared latent vector `z` * `trainer.G.mask(z=z)` to infer the mask * `trainer.compute_fire(x, segmentation)` to create a wildfire image from `x` and inferred segmentation * `trainer.compute_smog(x, depth)` to create a smog image from `x` and inferred depth * `trainer.compute_flood(x, mask)` to create a flood image from `x` and inferred mask using the painter (`trainer.G.paint(m, x)`) * `Trainer.resume_from_path()` static method to resume a trainer from a path ### Discriminator ## updates multi-batch: ``` multi_domain_batch = {"rf: batch0, "r": batch1, "s": batch2} ``` ## interfaces ### batches ```python batch = Dict({ "data": { "d": depthmap,, "s": segmentation_map, "m": binary_mask "x": real_flooded_image, }, "paths":{ same_keys: path_to_file } "domain": list(rf | r | s), "mode": list(train | val) }) ``` ### data #### json files | name | domain | description | author | | :--------------------------------------------- | :----: | :------------------------------------------------------------------------- | :-------: | | **train_r_full.json, val_r_full.json** | r | MiDaS+ Segmentation pseudo-labels .pt (HRNet + Cityscapes) | Mélisande | | **train_s_full.json, val_s_full.json** | s | Simulated data from Unity11k urban + Unity suburban dataset | *** | | train_s_nofences.json, val_s_nofences.json | s | Simulated data from Unity11k urban + Unity suburban dataset without fences | Alexia | | train_r_full_pl.json, val_r_full_pl.json | r | MegaDepth + Segmentation pseudo-labels .pt (HRNet + Cityscapes) | Alexia | | train_r_full_midas.json, val_r_full_midas.json | r | MiDaS+ Segmentation (HRNet + Cityscapes) | Mélisande | | train_r_full_old.json, val_r_full_old.json | r | MegaDepth+ Segmentation (HRNet + Cityscapes) | *** | | train_r_nopeople.json, val_r_nopeople.json | r | Same training data as above with people removed | Sasha | | train_rf_with_sim.json | rf | Doubled train_rf's size with sim data (randomly chosen) | Victor | | train_rf.json | rf | UPDATE (12/12/20): added 50 ims & masks from ADE20K Outdoors | Victor | | train_allres.json, val_allres.json | rf | includes both lowres and highres from ORCA_water_seg | Tianyu | | train_highres_only.json, val_highres_only.json | rf | includes only highres from ORCA_water_seg | Tianyu | ```yaml # data file ; one for each r|s - x: /path/to/image m: /path/to/mask s: /path/to/segmentation map - x: /path/to/another image d: /path/to/depth map m: /path/to/mask s: /path/to/segmentation map - x: ... ``` or ```json [ { "x": "/Users/victor/Documents/ccai/github/climategan/example_data/gsv_000005.jpg", "s": "/Users/victor/Documents/ccai/github/climategan/example_data/gsv_000005.npy", "d": "/Users/victor/Documents/ccai/github/climategan/example_data/gsv_000005_depth.jpg" }, { "x": "/Users/victor/Documents/ccai/github/climategan/example_data/gsv_000006.jpg", "s": "/Users/victor/Documents/ccai/github/climategan/example_data/gsv_000006.npy", "d": "/Users/victor/Documents/ccai/github/climategan/example_data/gsv_000006_depth.jpg" } ] ``` The json files used are located at `/network/tmp1/ccai/data/climategan/`. In the basenames, `_s` denotes simulated domain data and `_r` real domain data. The `base` folder contains json files with paths to images (`"x"`key) and masks (taken as ground truth for the area that should be flooded, `"m"` key). The `seg` folder contains json files and keys `"x"`, `"m"` and `"s"` (segmentation) for each image. loaders ``` loaders = Dict({ train: { r: loader, s: loader}, val: { r: loader, s: loader} }) ``` ### losses `trainer.losses` is a dictionary mapping to loss functions to optimize for the 3 main parts of the architecture: generator `G`, discriminators `D`: ```python trainer.losses = { "G":{ # generator "gan": { # gan loss from the discriminators "a": GANLoss, # adaptation decoder "t": GANLoss # translation decoder }, "cycle": { # cycle-consistency loss "a": l1 | l2,, "t": l1 | l2, }, "auto": { # auto-encoding loss a.k.a. reconstruction loss "a": l1 | l2, "t": l1 | l2 }, "tasks": { # specific losses for each auxillary task "d": func, # depth estimation "h": func, # height estimation "s": cross_entropy_2d, # segmentation "w": func, # water generation }, "classifier": l1 | l2 | CE # loss from fooling the classifier }, "D": GANLoss, # discriminator losses from the generator and true data "C": l1 | l2 | CE # classifier should predict the right 1-h vector [rf, rn, sf, sn] } ``` ## Logging on comet Comet.ml will look for api keys in the following order: argument to the `Experiment(api_key=...)` call, `COMET_API_KEY` environment variable, `.comet.config` file in the current working directory, `.comet.config` in the current user's home directory. If your not managing several comet accounts at the same time, I recommend putting `.comet.config` in your home as such: ``` [comet] api_key= workspace=vict0rsch rest_api_key= ``` ### Tests Run tests by executing `python test_trainer.py`. You can add `--no_delete` not to delete the comet experiment at exit and inspect uploads. Write tests as scenarios by adding to the list `test_scenarios` in the file. A scenario is a dict of overrides over the base opts in `shared/trainer/defaults.yaml`. You can create special flags for the scenario by adding keys which start with `__`. For instance, `__doc` is a mandatory key in any scenario describing it succinctly. ## Resources [Tricks and Tips for Training a GAN](https://chloes-dl.com/2019/11/19/tricks-and-tips-for-training-a-gan/) [GAN Hacks](https://github.com/soumith/ganhacks) [Keep Calm and train a GAN. Pitfalls and Tips on training Generative Adversarial Networks](https://medium.com/@utk.is.here/keep-calm-and-train-a-gan-pitfalls-and-tips-on-training-generative-adversarial-networks-edd529764aa9) ## Example **Inference: computing floods** ```python from pathlib import Path from skimage.io import imsave from tqdm import tqdm from climategan.trainer import Trainer from climategan.utils import find_images from climategan.tutils import tensor_ims_to_np_uint8s from climategan.transforms import PrepareInference model_path = "some/path/to/output/folder" # not .ckpt input_folder = "path/to/a/folder/with/images" output_path = "path/where/images/will/be/written" # resume trainer trainer = Trainer.resume_from_path(model_path, new_exp=None, inference=True) # find paths for all images in the input folder. There is a recursive option. im_paths = sorted(find_images(input_folder), key=lambda x: x.name) # Load images into tensors # * smaller side resized to 640 - keeping aspect ratio # * then longer side is cropped in the center # * result is a 1x3x640x640 float tensor in [-1; 1] xs = PrepareInference()(im_paths) # send to device xs = [x.to(trainer.device) for x in xs] # compute flood # * compute mask # * binarize mask if bin_value > 0 # * paint x using this mask ys = [trainer.compute_flood(x, bin_value=0.5) for x in tqdm(xs)] # convert 1x3x640x640 float tensors in [-1; 1] into 640x640x3 numpy arrays in [0, 255] np_ys = [tensor_ims_to_np_uint8s(y) for y in tqdm(ys)] # write images for i, n in tqdm(zip(im_paths, np_ys), total=len(im_paths)): imsave(Path(output_path) / i.name, n) ``` ## Release process In the `release/` folder * create a `model/` folder * create folders `model/masker/` and `model/painter/` * add the climategan code in `release/`: `git clone git@github.com:cc-ai/climategan.git` * move the code to `release/`: `cp climategan/* . && rm -rf climategan` * update `model/masker/opts/events` with `events:` from `shared/trainer/opts.yaml` * update `model/masker/opts/val.val_painter` to `"model/painter/checkpoints/latest_ckpt.pth"` * update `model/masker/opts/load_paths.m` to `"model/masker/checkpoints/latest_ckpt.pth"`