Model Training

Train a model for classifying pixels into different glacier types, using the pytorch package.

Code Link

1. In these notes, we’ll train a segmentation model to predict glacier, given the preprocessed patches created by the previous notes. Since we didn’t have enough time to preprocess all the training data in the previous notebook, we use the code block below to download and extract patches precomputed in advance. Note that there are only 67 patches in this archive – this keeps the download relatively quick, and keeps us from getting in trouble with the people hosting these binder notebooks.

import urllib.request
import tarfile
from pathlib import Path
from data import create_dir, download_data
import os

# setup directory structure for download
data_dir = Path("/home/jovyan/data")
process_dir = data_dir / "processed"
create_dir(process_dir)

# download processed data
download_data(
    "https://uwmadison.box.com/shared/static/d54agxzb5g8ivr7hkac8nygqd6nrgrqr.gz", 
    process_dir / "train.tar.gz"
)

2. The block below specifies some parameters of our learning algorithm, a U-Net segmentation model. Like most deep learning algorithms, we train it using a variant of stochastic optimization. The lr parameter below refers to the optimizer’s learning rate. The binder notebooks we’re running off of don’t have GPUs. If they did, we could set device: "cuda", and we’d be able to train the model much faster. We also had to limit the batch size, to avoid going over the memory limit imposed on these online notebooks.

args = {
    "batch_size": 1, # make this bigger if you are not running on binder
    "epochs": 50,
    "lr": 0.0001,
    "device": "cpu" # set to "cuda" if GPU is available
}

3. Our optimizer is going to need a way to stream in the preprocessed patches. This is accomplished using the DataLoader object below. If you tried visualizing the items in the data loader, you would see the same image-label pairs from the previous notebook. This step might seem mysterious if you haven’t used a deep learning algorithm before. I’m deliberately avoiding an extended discussion on deep learning – my emphasis here is on visualization and earth observation. There are also many good references on applied deep learning already.

from data import GlacierDataset
from torch.utils.data import DataLoader

paths = {
    "x": list((process_dir / "train").glob("x*")),
    "y": list((process_dir / "train").glob("y*"))
}

ds = GlacierDataset(paths["x"], paths["y"])
loader = DataLoader(ds, batch_size=args["batch_size"], shuffle=True)

4. Given a way of loading the training data, we can train our model. The parameters in the definition of the Unet correspond to 13 input sensor channels, 3 output classes (clean-ice glacier, debris-covered glacier, and background), and 4 layers.

5. We can try running the model below, but it will not finish in the time for this workshop (Though, with a larger batch size and a GPU, it doesn’t take too long to converge.). We’ll instead download a model that I already trained earlier. This model was also trained using all the training patches, and not just those from the Kokcha basin.

import torch.optim
from unet import Unet
from train import train_epoch

model = Unet(13, 3, 4, dropout=0.2).to(args["device"])
optimizer = torch.optim.Adam(model.parameters(), lr=args["lr"])

for epoch in range(args["epochs"]):
    train_epoch(model, loader, optimizer, args["device"], epoch)
    
torch.save(model.state_dict(), data_dir / "model.pt")