Training included SpaceNet models with the solaris Python API

We’ve included a number of SpaceNet models with solaris, including pre-trained model weights. You can find more information about your model choices here and the original competitors’ code for the models here.

For this tutorial we’ll walk through training a model using XD_XD’s SpaceNet 4 model. We’ll use the config file for that model, which you can find here.

You’ll also need to create training masks and create the image reference files before you start.

Once you’ve completed those steps, you can get down to model training! First, let’s load in the configuration:

[5]:
import solaris as sol

config = sol.utils.config.parse('/path/to/xdxd_spacenet4.yml')
config
[5]:
{'model_name': 'xdxd_spacenet4',
 'model_path': None,
 'train': False,
 'infer': True,
 'pretrained': True,
 'nn_framework': 'torch',
 'batch_size': 12,
 'data_specs': {'width': 512,
  'height': 512,
  'image_type': 'zscore',
  'rescale': False,
  'rescale_minima': 'auto',
  'rescale_maxima': 'auto',
  'channels': 4,
  'label_type': 'mask',
  'is_categorical': False,
  'mask_channels': 1,
  'val_holdout_frac': 0.2,
  'data_workers': None},
 'training_data_csv': '/path/to/training_df.csv',
 'validation_data_csv': None,
 'inference_data_csv': '/path/to/test_df.csv',
 'training_augmentation': {'augmentations': {'DropChannel': {'idx': 3,
    'axis': 2},
   'HorizontalFlip': {'p': 0.5},
   'RandomRotate90': {'p': 0.5},
   'RandomCrop': {'height': 512, 'width': 512, 'p': 1.0},
   'Normalize': {'mean': [0.006479, 0.009328, 0.01123],
    'std': [0.004986, 0.004964, 0.00495],
    'max_pixel_value': 65535.0,
    'p': 1.0}},
  'p': 1.0,
  'shuffle': True},
 'validation_augmentation': {'augmentations': {'DropChannel': {'idx': 3,
    'axis': 2},
   'CenterCrop': {'height': 512, 'width': 512, 'p': 1.0},
   'Normalize': {'mean': [0.006479, 0.009328, 0.01123],
    'std': [0.004986, 0.004964, 0.00495],
    'max_pixel_value': 65535.0,
    'p': 1.0}},
  'p': 1.0},
 'inference_augmentation': {'augmentations': {'DropChannel': {'idx': 3,
    'axis': 2,
    'p': 1.0},
   'Normalize': {'mean': [0.006479, 0.009328, 0.01123],
    'std': [0.004986, 0.004964, 0.00495],
    'max_pixel_value': 65535.0,
    'p': 1.0}},
  'p': 1.0},
 'training': {'epochs': 60,
  'steps_per_epoch': None,
  'optimizer': 'Adam',
  'lr': 0.0001,
  'opt_args': None,
  'loss': {'bcewithlogits': None, 'jaccard': None},
  'loss_weights': {'bcewithlogits': 10, 'jaccard': 2.5},
  'metrics': {'training': None, 'validation': None},
  'checkpoint_frequency': 10,
  'callbacks': {'model_checkpoint': {'filepath': 'xdxd_best.pth',
    'monitor': 'val_loss'}},
  'model_dest_path': 'xdxd.pth',
  'verbose': True},
 'inference': {'window_step_size_x': None,
  'window_step_size_y': None,
  'output_dir': 'inference_out/'}}

As you can see, the YAML gets parsed into a set of nested dictionaries by solaris. Relevant pieces of that config then get read during training.

Let’s assume that all of the paths in that config are correct. Now, to run training, you can run the following line:

[ ]:
trainer = sol.nets.train.Trainer(config)
trainer.train()

That’s all there is to it! With just four lines of Python code, you can train a model for geospatial ML!

If you wish to use a custom ML model not provided by solaris, check out the custom model training tutorial.