Repository: microsoft/torchgeo

Issue:
Benchmarking of GeoDataset for a paper result
### Datasets

We want to test several popular image sources, as well as both raster and vector labels.

- [ ] NAIP + Chesapeake
- [ ] Landsat + CDL
- [ ] Sentinel + Canadian Building Footprints

There is also a question of which file formats to test. For example, sampling from GeoJSON can take 3 min per getitem, whereas ESRI Shapefile only takes 1 sec per getitem (https://github.com/microsoft/torchgeo/pull/69#issuecomment-892816306).

### Experiments

- [ ] **Data location:** local vs. blob storage
- [ ] **Sampling strategy:** random vs. 2-step random tile/chip vs. non-random grid sampling
- [ ] **Warping strategy:** on-the-fly vs. pre-processing step
- [ ] **I/O strategy:** load/warp entire file vs. load/warp/cache entire file vs. load/warp single window

For the warping strategy, we should test the following possibilities:

- [ ] Already in correct CRS/resolution
- [ ] Need to change CRS
- [ ] Need to change resolution
- [ ] Need to change both CRS and resolution

What is the upfront cost of these pre-processing steps?

Example notebook: https://gist.github.com/calebrob6/d9bc5609ff638d601e2c35a1ab0a2dec


Discussion:
I think this will require a significant rework of our `__getitem__` implementation. Right now, we warp and then merge/sample from a tile at the same time. If we want to benefit from the 2-step random tile/chip sampling strategy, we'll have to use an LRU cache on the entire tile after warping.
I think we can also consider the following I/O strategies:

- [ ] load/warp entire file, no caching (worst case scenario)
- [ ] load/warp entire file, caching (good default)
- [ ] load/warp single window (does not allow for caching)

Merging should happen after the fact so that (tile 1, tile 2, tile 1 + 2) don't end up being 3 different entries in the cache.

I don't think we need to consider situations in which we:

* warp the file on disk on-the-fly (violates principle of least surprise)
* don't bother to warp (won't work if not already in correct CRS/res)
* don't bother to merge (won't be correct near boundaries)
* always return entire tiles instead of chips (not feasible for model/GPU memory)
* load the entire dataset into memory beforehand (won't fit in RAM)

These strategies make sense for tile-based raster images, but are slightly more complicated for vector geometries or static regional maps. We may need to change the default behavior based on the dataset.
For timing, we should choose some arbitrary epoch size, then experiment with various batch sizes and see how long it takes to load an entire epoch.
Here's where I'm currently stuck to remind myself when I next pick this up:

Our process right now is:

1. Open filehandles for raw data (`rasterio.open`)
2. Open filehandles for warped VRTs (`rasterio.vrt.WarpedVRT`)
3. Merge VRTs to get an array (`rasterio.merge.merge`)
4. Return array as tensor

Steps 1 and 2 don't actually do anything and are almost instantaneous. It isn't until you actually try to `read()` the data that warping occurs, and `read()` is called in `rasterio.merge.merge`. If we want to cache this reading of warped data, we'll have to call `vrt.read()` ourselves. Since `rasterio.merge.merge` only accepts filenames or filehandles as input, we'll basically need to implement our own `merge` algorithm that takes 1+ cached numpy arrays, creates a new array with the correct dimensions, and indexes the old arrays to copy the data. The hard part here will be keeping track of coordinates, nodata values, and merging correctly. See https://github.com/mapbox/rasterio/blob/master/rasterio/merge.py for the source code, most of which we'll need to do as well.
Another hurdle: the size of each array depends greatly on the dataset, but most are around 0.5 GB per file. We can't really assume users have >8 GB of RAM, which greatly limits our LRU cache size. We could use something like `psutil` to query the system memory, and hard-code the avg file size for each dataset if we want to make things more flexible.
For now, I think we can rely on GDAL's internal caching behavior. When I read a VRT the second time around, it seems to be significantly faster. Still not as fast as reading the raw data or as indexing from a loaded array, but good enough for a first round of benchmarking. GDAL also lets you configure the cache size.
Preliminary results look very promising!
![benchmark](https://user-images.githubusercontent.com/12021217/129969052-98987960-9519-4523-bfeb-b43c048384a9.png)


@adamjstewart, sketch of the full experiment:
- Get Landsat scenes from several projections and CDL data.
  - Convert all to COG if they aren't already
- Copy the files to blob storage, local SSD, local HDD
- Create three GeoDataset instances, one for each dataset location
- Record the number of patches per second you can read using each of the GeoDatasets with the different types of GeoSamplers
- Record how long it takes to warp/reproject the Landsat scenes to align them to CDL (or vice-versa) "by hand" with gdalwarp. 
  - Also record the size of the resulting files.
  - Also record the nasty gdalwarp command you actually have to figure out and execute to do this.
  - Note: We can use this to extrapolate how much preprocessing you would need to do before training with a traditional DL library.
  - Note: We can use this pre-aligned data with a custom dataloader to see how many patches/second you could sample if you did go in and do all the preprocessing. Hopefully this number is similar to what you get with torchgeo (or at least not much much larger).
@calebrob6 the above proposal covers the matrix of:

* **Data location:** local SSD, local HDD, blob storage
* **Sampling strategy:** RandomGeoSampler, RandomBatchGeoSampler, GridGeoSampler
* **I/O strategy:** cached, not cached

There are a lot of additional constraints that we're currently skipping:

* **File format:** GeoTIFF vs. HDF5, Shapefile vs. GeoJSON
* **Warping strategy:** already in correct CRS/res, change CRS, change res, change CRS and res

Do you think it's fine to skip these for the sake of time? I doubt reviewers would straight up reject us for not including one of these permutations, and can always ask us to perform additional experiments if they want.

Also, we should definitely benchmark not only `RasterDataset` but also `VectorDataset` (maybe Sentinel + Canadian Building Footprints?). Should I purposefully change the resolution of one of these datasets? Should I purposefully switch to a CRS different than all files or keep the CRS of one of the files?
Also, do we want to compare with different batch_sizes or different num_workers?
I'd do the first matrix as quickly as possible because the results of that are going to be very informative. If that all works out then you can repeat the same with a vectordataset.

> File format: GeoTIFF vs. HDF5, Shapefile vs. GeoJSON

I don't think this is important right now. I.e. we can just assume the data is in a good format (COG and shapefile/geopackage)

> Warping strategy

In the above sketch you can repeat the experiments with the manually aligned versions of the dataset to test the "already in correct CRS/res" case. The first set of experiments is with "change CRS and res". It might be interesting to see if warping or resampling is more expensive, but not interesting for the paper I think.

> Also, do we want to compare with different batch_sizes or different num_workers?

Sure! These experiments should be very quick to run once you have a script for them.

Resolving patch:
diff --git a/benchmark.py b/benchmark.py
new file mode 100755
--- /dev/null
+++ b/benchmark.py
@@ -0,0 +1,310 @@
+#!/usr/bin/env python3
+
+"""dataset and sampler benchmarking script."""
+
+import argparse
+import csv
+import os
+import time
+
+import pytorch_lightning as pl
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader
+from torchvision.models import resnet18
+
+from torchgeo.datasets import CDL, BoundingBox, Landsat8
+from torchgeo.samplers import GridGeoSampler, RandomBatchGeoSampler, RandomGeoSampler
+
+
+def set_up_parser() -> argparse.ArgumentParser:
+    """Set up the argument parser.
+
+    Returns:
+        the argument parser
+    """
+    parser = argparse.ArgumentParser(
+        description=__doc__, formatter_class=argparse.ArgumentDefaultsHelpFormatter
+    )
+
+    parser.add_argument(
+        "--landsat-root",
+        default=os.path.join("data", "landsat"),
+        help="directory containing Landsat data",
+        metavar="ROOT",
+    )
+    parser.add_argument(
+        "--cdl-root",
+        default=os.path.join("data", "cdl"),
+        help="directory containing CDL data",
+        metavar="ROOT",
+    )
+    parser.add_argument(
+        "-d", "--device", default=0, type=int, help="CPU/GPU ID to use", metavar="ID"
+    )
+    parser.add_argument(
+        "-c",
+        "--cache",
+        action="store_true",
+        help="cache file handles during data loading",
+    )
+    parser.add_argument(
+        "-b",
+        "--batch-size",
+        default=2 ** 4,
+        type=int,
+        help="number of samples in each mini-batch",
+        metavar="SIZE",
+    )
+    group = parser.add_mutually_exclusive_group(required=True)
+    group.add_argument(
+        "-n",
+        "--num-batches",
+        type=int,
+        help="number of batches to load",
+        metavar="SIZE",
+    )
+    group.add_argument(
+        "-e",
+        "--epoch-size",
+        type=int,
+        help="number of samples to load, should be evenly divisible by batch size",
+        metavar="SIZE",
+    )
+    parser.add_argument(
+        "-p",
+        "--patch-size",
+        default=224,
+        type=int,
+        help="height/width of each patch",
+        metavar="SIZE",
+    )
+    parser.add_argument(
+        "-s",
+        "--stride",
+        default=2 ** 7,
+        type=int,
+        help="sampling stride for GridGeoSampler",
+    )
+    parser.add_argument(
+        "-w",
+        "--num-workers",
+        default=0,
+        type=int,
+        help="number of workers for parallel data loading",
+        metavar="NUM",
+    )
+    parser.add_argument(
+        "--seed",
+        default=0,
+        type=int,
+        help="random seed for reproducibility",
+    )
+    parser.add_argument(
+        "--output-fn",
+        default="benchmark-results.csv",
+        type=str,
+        help="path to the CSV file to write results",
+        metavar="FILE",
+    )
+    parser.add_argument(
+        "-v",
+        "--verbose",
+        action="store_true",
+        help="print results to stdout",
+    )
+
+    return parser
+
+
+def main(args: argparse.Namespace) -> None:
+    """High-level pipeline.
+
+    Benchmarks performance of various samplers with and without caching.
+
+    Args:
+        args: command-line arguments
+    """
+    bands = ["B1", "B2", "B3", "B4", "B5", "B6", "B7"]
+
+    # Benchmark samplers
+
+    # Initialize datasets
+    cdl = CDL(args.cdl_root, cache=args.cache)
+    landsat = Landsat8(
+        args.landsat_root, crs=cdl.crs, res=cdl.res, cache=args.cache, bands=bands
+    )
+    dataset = landsat + cdl
+
+    # Initialize samplers
+    if args.epoch_size:
+        length = args.epoch_size
+        num_batches = args.epoch_size // args.batch_size
+    elif args.num_batches:
+        length = args.num_batches * args.batch_size
+        num_batches = args.num_batches
+
+    # Workaround for https://github.com/microsoft/torchgeo/issues/149
+    roi = BoundingBox(
+        -2000000, 2200000, 280000, 3170000, dataset.bounds.mint, dataset.bounds.maxt
+    )
+    samplers = [
+        RandomGeoSampler(
+            landsat.index,
+            size=args.patch_size,
+            length=length,
+            roi=roi,
+        ),
+        GridGeoSampler(
+            landsat.index, size=args.patch_size, stride=args.stride, roi=roi
+        ),
+        RandomBatchGeoSampler(
+            landsat.index,
+            size=args.patch_size,
+            batch_size=args.batch_size,
+            length=length,
+            roi=roi,
+        ),
+    ]
+
+    results_rows = []
+    for sampler in samplers:
+        if args.verbose:
+            print(f"\n{sampler.__class__.__name__}:")
+
+        if isinstance(sampler, RandomBatchGeoSampler):
+            dataloader = DataLoader(
+                dataset,
+                batch_sampler=sampler,  # type: ignore[arg-type]
+                num_workers=args.num_workers,
+            )
+        else:
+            dataloader = DataLoader(
+                dataset,
+                batch_size=args.batch_size,
+                sampler=sampler,  # type: ignore[arg-type]
+                num_workers=args.num_workers,
+            )
+
+        tic = time.time()
+        num_total_patches = 0
+        for i, batch in enumerate(dataloader):
+            num_total_patches += args.batch_size
+            # This is to stop the GridGeoSampler from enumerating everything
+            if i == num_batches - 1:
+                break
+        toc = time.time()
+        duration = toc - tic
+
+        if args.verbose:
+            print(f"  duration: {duration:.3f} sec")
+            print(f"  count: {num_total_patches} patches")
+            print(f"  rate: {num_total_patches / duration:.3f} patches/sec")
+
+        if args.cache:
+            if args.verbose:
+                print(landsat._cached_load_warp_file.cache_info())
+
+            # Clear cache for fair comparison between samplers
+            # Both `landsat` and `cdl` share the same cache
+            landsat._cached_load_warp_file.cache_clear()
+
+        results_rows.append(
+            {
+                "cached": args.cache,
+                "seed": args.seed,
+                "duration": duration,
+                "count": num_total_patches,
+                "rate": num_total_patches / duration,
+                "sampler": sampler.__class__.__name__,
+                "batch_size": args.batch_size,
+                "num_workers": args.num_workers,
+            }
+        )
+
+    # Benchmark model
+    model = resnet18()
+    # Change number of input channels to match Landsat
+    model.conv1 = nn.Conv2d(  # type: ignore[attr-defined]
+        len(bands), 64, kernel_size=7, stride=2, padding=3, bias=False
+    )
+
+    criterion = nn.CrossEntropyLoss()  # type: ignore[attr-defined]
+    params = model.parameters()
+    optimizer = optim.SGD(params, lr=0.0001)
+
+    device = torch.device(  # type: ignore[attr-defined]
+        "cuda" if torch.cuda.is_available() else "cpu", args.device
+    )
+    model = model.to(device)
+
+    tic = time.time()
+    num_total_patches = 0
+    for _ in range(num_batches):
+        num_total_patches += args.batch_size
+        x = torch.rand(args.batch_size, len(bands), args.patch_size, args.patch_size)
+        # y = torch.randint(0, 256, (args.batch_size, args.patch_size, args.patch_size))
+        y = torch.randint(0, 256, (args.batch_size,))  # type: ignore[attr-defined]
+        x = x.to(device)
+        y = y.to(device)
+
+        optimizer.zero_grad()
+        prediction = model(x)
+        loss = criterion(prediction, y)
+        loss.backward()
+        optimizer.step()
+
+    toc = time.time()
+    duration = toc - tic
+
+    if args.verbose:
+        print("\nResNet-18:")
+        print(f"  duration: {duration:.3f} sec")
+        print(f"  count: {num_total_patches} patches")
+        print(f"  rate: {num_total_patches / duration:.3f} patches/sec")
+
+    results_rows.append(
+        {
+            "cached": args.cache,
+            "seed": args.seed,
+            "duration": duration,
+            "count": num_total_patches,
+            "rate": num_total_patches / duration,
+            "sampler": "resnet18",
+            "batch_size": args.batch_size,
+            "num_workers": args.num_workers,
+        }
+    )
+
+    fieldnames = [
+        "cached",
+        "seed",
+        "duration",
+        "count",
+        "rate",
+        "sampler",
+        "batch_size",
+        "num_workers",
+    ]
+    if not os.path.exists(args.output_fn):
+        with open(args.output_fn, "w") as f:
+            writer = csv.DictWriter(f, fieldnames=fieldnames)
+            writer.writeheader()
+    with open(args.output_fn, "a") as f:
+        writer = csv.DictWriter(f, fieldnames=fieldnames)
+        writer.writerows(results_rows)
+
+
+if __name__ == "__main__":
+    os.environ["GDAL_CACHEMAX"] = "50%"
+
+    parser = set_up_parser()
+    args = parser.parse_args()
+
+    if args.epoch_size:
+        assert args.epoch_size % args.batch_size == 0
+
+    pl.seed_everything(args.seed)
+
+    main(args)
diff --git a/run_benchmarks_experiments.py b/run_benchmarks_experiments.py
new file mode 100755
--- /dev/null
+++ b/run_benchmarks_experiments.py
@@ -0,0 +1,44 @@
+"""Script for running the benchmark script over a sweep of different options."""
+import itertools
+import subprocess
+import time
+from typing import List
+
+NUM_BATCHES = 100
+
+SEED_OPTIONS = [0, 1, 2]
+CACHE_OPTIONS = [True, False]
+BATCH_SIZE_OPTIONS = [16, 32, 64, 128, 256, 512]
+
+total_num_experiments = len(SEED_OPTIONS) * len(CACHE_OPTIONS) * len(BATCH_SIZE_OPTIONS)
+
+if __name__ == "__main__":
+
+    tic = time.time()
+    for i, (cache, batch_size, seed) in enumerate(
+        itertools.product(CACHE_OPTIONS, BATCH_SIZE_OPTIONS, SEED_OPTIONS)
+    ):
+        print(f"{i}/{total_num_experiments} -- {time.time() - tic}")
+        tic = time.time()
+        command: List[str] = [
+            "python",
+            "benchmark.py",
+            "--landsat-root",
+            "/datadrive/landsat",
+            "--cdl-root",
+            "/datadrive/cdl",
+            "-w",
+            "6",
+            "--batch-size",
+            str(batch_size),
+            "--num-batches",
+            str(NUM_BATCHES),
+            "--seed",
+            str(seed),
+            "--verbose",
+        ]
+
+        if cache:
+            command.append("--cache")
+
+        subprocess.call(command)


Test patch: