UNET Training on Multi-Class Segmentation from Satellite Imagery (DSTL)

• Hardware
-Tesla V100
-Number of GPU’s used = 4
• Network Type
• TLT Version (Please run “tlt info --verbose” and share “docker_tag” here)
-TAO 3.21.11
• Training spec file(If have, please share here)

We are trying to train UNET on this data that we chop into smaller size (572 and 256):
-Data: Dstl Satellite Imagery Feature Detection | Kaggle

Attached is:

UNETResNet_experiment_spec.txt (spec file for UNET training with ResNet 50 backbone and 256x256 network size).

UNRN_image.png (training images from DSTL link, above, chopped into 256x256 size)

UNRN_label.png (label for the image)

UNRN_inference.png (a resulting inference after training)

resnet_50_results_tlt.json (the F1 scores per class)

Vanallia_UNET_experiment_spec.txt (spec file for UNET training with Vanalla UNET backbone, and 572 x572 network size)

VU_original.png (image of 572x572 size)

VU_label.png (label of 572x572 for the image)

VU_inferenceLabel.png (resulting image from inference after training)

VU_labeledImage.png (Ground truth over image to show the less than ideal results)

vanilla_unet_results_tlt.json (the F1 scores per class)

ColorMapping.txt (showing which colors are for which classes)

I can provide more examples if you need them.

We are in need of a good architecture, and good hyper parameters for segmenting satellite data. I understand that UNET was originally designed for a one class problem for medical imaging, but I have seen successful use-cases on multi-class data. I have tried the cross-entropy loss (which is suggested for multi-class), and other variations.

Can you help us get better results with UNET on this data?

I appreciate your time and help,


UNTEResNet_experiment_spec.txt (2.6 KB)
VanillaUNET_experiment_spec.txt (2.6 KB)

ColorMapping.txt (211 Bytes)
resnet_50_results_tlt.json (567 Bytes)
vanilla_unet_results_tlt.json (600 Bytes)

For improving accuracy, refer to Problems encountered in training unet and inference unet - #27 by Morganh , you can also use

  • loss: “cross_entropy”
  • weight: 2e-06
  • crop_and_resize_prob : 0.01

Howdy @Morganh,

Thank you for the suggestion. I’m working with @derek.r.skilling on this project. We ran with your suggested configuration on a ResNet18 architecture and got the following results:

  • Recall : 0.39845550229603593
  • Precision: 0.6882064789533615
  • F1 score: 0.5690247036853558
  • Mean IOU: 0.32365456951612775

We previously also ran with Vanilla UNET but with our previous configuration(See VanillaUNET_experiment_spec.txt file in first post) and got the following results:

  • Recall : 0.6660359132484618
  • Precision: 0.8501717475327578
  • F1 score: 0.6888677606271899
  • Mean IOU: 0.6156709645239805

Considering that the Vanilla UNET architecture evaluation scores, are there any suggestions for getting better results? Either with architecture or any of the available UNET architectures?

Thank you for your time

The latest TAO 22.05 version provides below architectures.

["resnet", "vgg", "vanilla_unet", "efficientnet_b0",
  "vanilla_unet_dynamic", "shufflenet", "byom"]

Please try to use vanilla_unet_dynamic and efficientnet_b0 .

For example,

model_config {
  num_layers: 18
  model_input_width: 960
  model_input_height: 544
  model_input_channels: 3
all_projections: true
arch: "vanilla_unet_dynamic"
use_batch_norm: true
enable_qat: true
training_precision {
backend_floatx: FLOAT32

training_config {
  batch_size: 1
  epochs: 100
  log_summary_steps: 100
  checkpoint_interval: 1
  loss: "cross_entropy"
  regularizer {
    type: L2
    weight: 1e-5
  optimizer {
    adam {
      epsilon: 9.99999993923e-09
      beta1: 0.899999976158
      beta2: 0.999000012875
  visualizer {
  enabled: false
  data_options: true
1 Like

Thank you for the suggestion @Morganh. Quick follow-up. We’re currently running on TAO 3.21.11 and unfortunately we’re working in an airgaped network. We see that we don’t have vanilla_unet_dynamic as an available architecture in the documentation. Is this similar to the vanilla_dynamic architecture that is available for the 3.21.11 release?

The 3.21.11 release supports

["resnet", "vgg", "vanilla_unet", "efficientnet_b0",

You can set vanilla_unet_dynamic.

1 Like

There is no update from you for a period, assuming this is not an issue anymore.
Hence we are closing this topic. If need further support, please open a new one.

This topic was automatically closed 14 days after the last reply. New replies are no longer allowed.