Prepare self-supervised imaginative and prescient transformers on overhead imagery with Amazon SageMaker

This can be a visitor weblog submit co-written with Ben Veasey, Jeremy Anderson, Jordan Knight, and June Li from Vacationers.

Satellite tv for pc and aerial photos present perception into a variety of issues, together with precision agriculture, insurance coverage threat evaluation, city improvement, and catastrophe response. Coaching machine studying (ML) fashions to interpret this information, nevertheless, is bottlenecked by expensive and time-consuming human annotation efforts. One strategy to overcome this problem is thru self-supervised learning (SSL). By coaching on massive quantities of unlabeled picture information, self-supervised fashions be taught picture representations that may be transferred to downstream duties, reminiscent of picture classification or segmentation. This method produces picture representations that generalize properly to unseen information and reduces the quantity of labeled information required to construct performant downstream fashions.

On this submit, we display the right way to practice self-supervised imaginative and prescient transformers on overhead imagery utilizing Amazon SageMaker. Vacationers collaborated with the Amazon Machine Studying Options Lab (now often called the Generative AI Innovation Center) to develop this framework to assist and improve aerial imagery mannequin use circumstances. Our resolution relies on the DINO algorithm and makes use of the SageMaker distributed data parallel library (SMDDP) to separate the info over a number of GPU cases. When pre-training is full, the DINO picture representations might be transferred to quite a lot of downstream duties. This initiative led to improved mannequin performances inside the Vacationers Information & Analytics house.

Overview of resolution

The 2-step course of for pre-training imaginative and prescient transformers and transferring them to supervised downstream duties is proven within the following diagram.

Within the following sections, we offer a walkthrough of the answer utilizing satellite tv for pc photos from the BigEarthNet-S2 dataset. We construct on the code offered within the DINO repository.

Stipulations

Earlier than getting began, you want entry to a SageMaker notebook instance and an Amazon Simple Storage Service (Amazon S3) bucket.

Put together the BigEarthNet-S2 dataset

BigEarthNet-S2 is a benchmark archive that incorporates 590,325 multispectral photos collected by the Sentinel-2 satellite tv for pc. The pictures doc the land cowl, or bodily floor options, of ten European international locations between June 2017 and Could 2018. The forms of land cowl in every picture, reminiscent of pastures or forests, are annotated in line with 19 labels. The next are just a few instance RGB photos and their labels.

Step one in our workflow is to arrange the BigEarthNet-S2 dataset for DINO coaching and analysis. We begin by downloading the dataset from the terminal of our SageMaker pocket book occasion:

wget https://bigearth.internet/downloads/BigEarthNet-S2-v1.0.tar.gz
tar -xvf BigEarthNet-S2-v1.0.tar.gz

The dataset has a dimension of about 109 GB. Every picture is saved in its personal folder and incorporates 12 spectral channels. Three bands with 60m spatial decision (60-meter pixel top/width) are designed to establish aerosols (B01), water vapor (B09), and clouds (B10). Six bands with 20m spatial decision are used to establish vegetation (B05, B06, B07, B8A) and distinguish between snow, ice, and clouds (B11, B12). Three bands with 10m spatial decision assist seize seen and near-infrared mild (B02, B03, B04, B8/B8A). Moreover, every folder incorporates a JSON file with the picture metadata. An in depth description of the info is offered within the BigEarthNet Guide.

To carry out statistical analyses of the info and cargo photos throughout DINO coaching, we course of the person metadata recordsdata into a standard geopandas Parquet file. This may be finished utilizing the BigEarthNet Widespread and the BigEarthNet GDF Builder helper packages:

python -m bigearthnet_gdf_builder.builder build-recommended-s2-parquet BigEarthNet-v1.0/

The ensuing metadata file incorporates the really useful picture set, which excludes 71,042 photos which might be absolutely coated by seasonal snow, clouds, and cloud shadows. It additionally incorporates data on the acquisition date, location, land cowl, and practice, validation, and take a look at cut up for every picture.

We retailer the BigEarthNet-S2 photos and metadata file in an S3 bucket. As a result of we use true colour photos throughout DINO coaching, we solely add the pink (B04), inexperienced (B03), and blue (B02) bands:

aws s3 cp final_ben_s2.parquet s3://bigearthnet-s2-dataset/metadata/
aws s3 cp BigEarthNet-v1.0/ s3://bigearthnet-s2-dataset/data_rgb/ 
    --recursive 
    --exclude "*" 
    --include "_B02.tif" 
    --include "_B03.tif"  
    --include "_B04.tif"

The dataset is roughly 48 GB in dimension and has the next construction:

bigearthnet-s2-dataset/                                    Amazon S3 bucket
├── metadata/
│ └── final_ben_s2.parquet 
└── dataset_rgb/
  ├── S2A_MSIL2A_20170613T101031_0_45/
  │ └── S2A_MSIL2A_20170613T101031_0_45_B02.tif            Blue channel
  │ └── S2A_MSIL2A_20170613T101031_0_45_B03.tif            Inexperienced channel
  │ └── S2A_MSIL2A_20170613T101031_0_45_B04.tif            Pink channel

Prepare DINO fashions with SageMaker

Now that our dataset has been uploaded to Amazon S3, we transfer to coach DINO fashions on BigEarthNet-S2. As proven within the following determine, the DINO algorithm passes totally different world and native crops of an enter picture to scholar and instructor networks. The scholar community is taught to match the output of the instructor community by minimizing the cross-entropy loss. The scholar and instructor weights are linked by an exponential shifting common (EMA).

We make two modifications to the unique DINO code. First, we create a customized PyTorch dataset class to load the BigEarthNet-S2 photos. The code was initially written to course of ImageNet information and expects photos to be saved by class. BigEarthNet-S2, nevertheless, is a multi-label dataset the place every picture resides in its personal subfolder. Our dataset class hundreds every picture utilizing the file path saved within the metadata:

import pandas as pd
import rasterio
from PIL import Picture
import torch
from torch.utils.information import Dataset, DataLoader
from torchvision import transforms, utils
 
OPTICAL_MAX_VALUE = 2000

LAND_COVER_LABELS = [
    "Urban fabric",
    "Industrial or commercial units",
    "Arable land",
    "Permanent crops",
    "Pastures",
    "Complex cultivation patterns",
    "Land principally occupied by agriculture, with significant areas of natural vegetation",
    "Agro-forestry areas",
    "Broad-leaved forest",
    "Coniferous forest",
    "Mixed forest",
    "Natural grassland and sparsely vegetated areas",
    "Moors, heathland and sclerophyllous vegetation",
    "Transitional woodland, shrub",
    "Beaches, dunes, sands",
    "Inland wetlands",
    "Coastal wetlands",
    "Inland waters",
    "Marine waters",
]
 
class BigEarthNetDataset(Dataset):
     """
     PyTorch dataset class that hundreds the BigEarthNet-S2 photos from a metadata file.

     Args: 
          metadata_file: path to metadata file 
          data_dir: listing the place BigEarthNet-S2 information is positioned  
          cut up: practice, validation, or take a look at cut up
          rework: transformations utilized to the enter picture
     """
     def __init__(self, metadata_file, data_dir, cut up="practice", rework=None):
		# picture file paths from metadata
        metadata = pd.read_parquet(metadata_file)
        self.metadata_split = metadata[metadata["original_split"] == cut up]
        self.data_dir = data_dir
        self.patch_names = self.metadata_split["name"].tolist()
 
        # one-hot-encode land cowl labels 
        multiclass_labels = self.metadata_split.new_labels.tolist()
        self.labels = self.get_multi_onehot_labels(multiclass_labels)

        # transforms        
        self.rework = rework
 
    def __len__(self):
        """Return size of dataset."""
        return len(self.metadata_split)
 
    def __getitem__(self, index):
        """Returns the picture and label for a given index."""
        patch_name = self.patch_names[index]
        file_path = os.path.be part of(self.data_dir, patch_name)
	
	# generate RGB picture
        r_channel = rasterio.open(os.path.be part of(file_path, patch_name + "_B04.tif")).learn(1)
        g_channel = rasterio.open(os.path.be part of(file_path, patch_name + "_B03.tif")).learn(1)
        b_channel = rasterio.open(os.path.be part of(file_path, patch_name + "_B02.tif")).learn(1)
 
        picture = np.stack([r_channel, g_channel, b_channel], axis=2)
        picture = picture / OPTICAL_MAX_VALUE * 255
        picture = np.clip(picture, 0, 225).astype(np.uint8)
    
        # apply picture transforms
        picture = Picture.fromarray(picture, mode="RGB")
        if self.rework is just not None:
            picture = self.rework(picture)
 
        # load label
        label = self.labels[index]
 
        return picture, label
  
    def get_multi_onehot_labels(self, multiclass_labels):
        """Convert BEN-19 labels to one-hot encoded vector."""
        targets = torch.zeros([len(multiclass_labels), len(LAND_COVER_LABELS)])
        for index, img_labels in enumerate(multiclass_labels):
            for label in img_labels:
                index_hot = LAND_COVER_LABELS.index(label)
                targets[index, index_hot] = 1.
        return targets

This dataset class is named in main_dino.py throughout coaching. Though the code features a perform to one-hot encode the land cowl labels, these labels are usually not utilized by the DINO algorithm.

The second change we make to the DINO code is so as to add assist for SMDDP. We add the next code to the init_distributed_mode perform within the util.py file:

init_distributed_mode perform within the util.py file:

def init_distributed_mode(args):
     if json.hundreds(
          os.environ.get('SM_FRAMEWORK_PARAMS', '{}'))
         .get('sagemaker_distributed_dataparallel_enabled', False)
     ): 
          # launch coaching with SMDDP 
          dist.init_process_group(backend='smddp')
          args.word_size = dist.get_world_size() 
          args.gpu = int(os.environ['LOCAL_RANK'])

With these changes, we’re prepared to coach DINO fashions on BigEarthNet-S2 utilizing SageMaker. To coach on a number of GPUs or cases, we create a SageMaker PyTorch Estimator that ingests the DINO coaching script, the picture and metadata file paths, and the coaching hyperparameters:

import time
from sagemaker.pytorch import PyTorch

# output bucket the place closing mannequin artifacts are uploaded 
DINO_OUTPUT_BUCKET = 'dino-models'

# paths on coaching occasion  
sm_metadata_path="/decide/ml/enter/information/metadata"              
sm_data_path="/decide/ml/enter/information/practice"                     
sm_output_path="/decide/ml/output/information"                        
sm_checkpoint_path="/decide/ml/checkpoints"                

# coaching job identify
dino_base_job_name = f'dino-model-{int(time.time())}'

# create SageMaker Estimator
estimator = PyTorch(
    base_job_name=dino_base_job_name,
    source_dir="path/to/aerial_featurizer",
    entry_point="main_dino.py",
    position=position,
    framework_version="1.12",
    py_version="py38",
    instance_count=1,
    instance_type="ml.p3.16xlarge",    
    distribution = {'smdistributed':{'dataparallel':{'enabled': True}}},        
    volume_size=100,
    sagemaker_session=sagemaker_session,
    hyperparameters = {
        # hyperparameters handed to entry level script
        'arch': 'vit_small',
        'patch_size': 16,
        'metadata_dir': sm_metadata_path,
        'data_dir': sm_data_path,
        'output_dir': sm_output_path,
        'checkpoint_dir': sm_checkpoint_path,
        'epochs': 100,
        'saveckp_freq': 20,
    },
    max_run=24*60*60,               
    checkpoint_local_path = sm_checkpoint_path,
    checkpoint_s3_uri =f's3://{DINO_OUTPUT_BUCKET}/checkpoints/{base_job_name}', 
    debugger_hook_config=False,                           
)

This code specifies that we’ll practice a small imaginative and prescient transformer mannequin (21 million parameters) with a patch dimension of 16 for 100 epochs. It’s best apply to create a brand new checkpoint_s3_uri for every coaching job with a view to cut back the preliminary information obtain time. As a result of we’re utilizing SMDDP, we should practice on an ml.p3.16xlarge, ml.p3dn.24xlarge, or ml.p4d.24xlarge occasion. It is because SMDDP is simply enabled for the biggest multi-GPU cases. To coach on smaller occasion sorts with out SMDDP, you will want to take away the distribution and debugger_hook_config arguments from the estimator.

After now we have created the SageMaker PyTorch Estimator, we launch the coaching job by calling the match technique. We specify the enter coaching information utilizing the Amazon S3 URIs for the BigEarthNet-S2 metadata and pictures:

# name match to start coaching
estimator.match(
    inputs={
        'metadata': 's3://bigearthnet-s2-dataset/metadata/',
        'practice': 's3://bigearthnet-s2-dataset/data_rgb/',
    },
    wait=False
)

SageMaker spins up the occasion, copies the coaching script and dependencies, and begins DINO coaching. We are able to monitor the progress of the coaching job from our Jupyter pocket book utilizing the next instructions:

# monitor coaching
training_job_name = estimator.latest_training_job.identify 
attached_estimator = PyTorch.connect(training_job_name)
attached_estimator.logs()

We are able to additionally monitor occasion metrics and examine log recordsdata on the SageMaker console below Coaching jobs. Within the following figures, we plot the GPU utilization and loss perform for a DINO mannequin skilled on an ml.p3.16xlarge occasion with a batch dimension of 128.

Throughout coaching, the GPU utilization is 83% of the ml.p3.16xlarge capability (8 NVIDIA Tesla V100 GPUs) and the VRAM utilization is 85%. The loss perform steadily decreases with every epoch, indicating that the outputs of the scholar and instructor networks have gotten extra comparable. In complete, coaching takes about 11 hours.

Switch studying to downstream duties

Our skilled DINO mannequin might be transferred to downstream duties like picture classification or segmentation. On this part, we use the pre-trained DINO options to foretell the land cowl lessons for photos within the BigEarthNet-S2 dataset. As depicted within the following diagram, we practice a multi-label linear classifier on high of frozen DINO options. On this instance, the enter picture is related to arable land and pasture land covers.

Many of the code for the linear classifier is already in place within the unique DINO repository. We make just a few changes for our particular process. As earlier than, we use the customized BigEarthNet dataset to load photos throughout coaching and analysis. The labels for the pictures are one-hot encoded as 19-dimensional binary vectors. We use the binary cross-entropy for the loss perform and compute the average precision to guage the efficiency of the mannequin.

To coach the classifier, we create a SageMaker PyTorch Estimator that runs the coaching script, eval_linear.py. The coaching hyperparameters embrace the main points of the DINO mannequin structure and the file path for the mannequin checkpoint:

# output bucket the place closing mannequin artifacts are uploaded 
CLASSIFIER_OUTPUT_BUCKET = 'land-cover-classification'

# DINO checkpoint identify 
checkpoint="checkpoint.pth"

# paths on coaching occasion  
sm_dino_path = f'/decide/ml/enter/information/dino_checkpoint'          
sm_dino_checkpoint = f'{sm_dino_path}/{checkpoint}'           

# coaching job identify
classifier_base_job_name = f'linear-classifier-{int(time.time())}'

# create Estimator 
estimator = PyTorch(
    base_job_name=classifier_base_job_name,
    source_dir="path/to/aerial_featurizer",
    entry_point="eval_linear.py",
    position=position,
    framework_version='1.12',
    py_version='py38',
    instance_count=1,
    instance_type="ml.p3.2xlarge",
    sagemaker_session=sagemaker_session,
    hyperparameters = {
    # hyperparameters handed to entry level script
        'arch': 'vit_small',
        'pretrained_weights': sm_dino_checkpoint,
        'epochs': 50,
        'data_dir': sm_data_path,
        'metadata_dir': sm_metadata_path,
        'output_dir': sm_checkpoint_path,
        'num_labels': 19,
    },
    max_run=1*60*60, 
    checkpoint_local_path = sm_checkpoint_path,
    checkpoint_s3_uri =f's3://{CLASSIFIER_OUTPUT_BUCKET}/checkpoints/{base_job_name}',
)

We begin the coaching job utilizing the match technique, supplying the Amazon S3 places of the BigEarthNet-S2 metadata and coaching photos and the DINO mannequin checkpoint:

# name match to start coaching
estimator.match(
    inputs={
    'metadata': 's3://bigearthnet-s2-dataset/metadata/',
    'dataset': 's3://bigearthnet-s2-dataset/data_rgb/',
    'dino_checkpoint': f's3://bigearthnet-s2-dataset/dino-models/checkpoints/{dino_base_job_name}',
    },
    wait=False
)

When coaching is full, we are able to carry out inference on the BigEarthNet-S2 take a look at set utilizing SageMaker batch transform or SageMaker Processing. Within the following desk, we examine the common precision of the linear mannequin on take a look at set photos utilizing two totally different DINO picture representations. The primary mannequin, ViT-S/16 (ImageNet), is the small imaginative and prescient transformer checkpoint included within the DINO repository that was pre-trained utilizing front-facing photos within the ImageNet dataset. The second mannequin, ViT-S/16 (BigEarthNet-S2), is the mannequin we produced by pre-training on overhead imagery.

We discover that the DINO mannequin pre-trained on BigEarthNet-S2 transfers higher to the land cowl classification process than the DINO mannequin pre-trained on ImageNet, leading to a 6.7% enhance within the common precision.

Clear up

After finishing DINO coaching and switch studying, we are able to clear up our sources to keep away from incurring costs. We stop or delete our notebook instance and remove any unwanted data or model artifacts from Amazon S3.

Conclusion

This submit demonstrated the right way to practice DINO fashions on overhead imagery utilizing SageMaker. We used SageMaker PyTorch Estimators and SMDDP with a view to generate representations of BigEarthNet-S2 photos with out the necessity for express labels. We then transferred the DINO options to a downstream picture classification process, which concerned predicting the land cowl class of BigEarthNet-S2 photos. For this process, pre-training on satellite tv for pc imagery yielded a 6.7% enhance in common precision relative to pre-training on ImageNet.

You need to use this resolution as a template for coaching DINO fashions on large-scale, unlabeled aerial and satellite tv for pc imagery datasets. To be taught extra about DINO and constructing fashions on SageMaker, take a look at the next sources:

In regards to the Authors

Ben Veasey is a Senior Affiliate Information Scientist at Vacationers, working inside the AI & Automation Accelerator staff. With a deep understanding of modern AI applied sciences, together with pc imaginative and prescient, pure language processing, and generative AI, Ben is devoted to accelerating the adoption of those applied sciences to optimize enterprise processes and drive effectivity at Vacationers.

Jeremy Anderson is a Director & Information Scientist at Vacationers on the AI & Automation Accelerator staff. He’s inquisitive about fixing enterprise issues with the most recent AI and deep studying strategies together with massive language fashions, foundational imagery fashions, and generative AI. Previous to Vacationers, Jeremy earned a PhD in Molecular Biophysics from the Johns Hopkins College and in addition studied evolutionary biochemistry. Outdoors of labor you will discover him operating, woodworking, or rewilding his yard.

Jordan Knight is a Senior Information Scientist working for Vacationers within the Enterprise Insurance coverage Analytics & Analysis Division. His ardour is for fixing difficult real-world pc imaginative and prescient issues and exploring new state-of-the-art strategies to take action. He has a specific curiosity within the social impression of ML fashions and the way we are able to proceed to enhance modeling processes to develop ML options which might be equitable for all. Jordan graduated from MIT with a Grasp’s in Enterprise Analytics. In his free time you will discover him both mountain climbing, mountaineering, or persevering with to develop his considerably rudimentary cooking expertise.

June Li is a knowledge scientist at Vacationers’s Enterprise Insurance coverage’s Synthetic Intelligence staff, the place she leads and coordinates work within the AI imagery portfolio. She is obsessed with implementing modern AI options that convey substantial worth to the enterprise companions and stakeholders. Her work has been integral in remodeling complicated enterprise challenges into alternatives by leveraging cutting-edge AI applied sciences.

Sourav Bhabesh is a Senior Utilized Scientist on the AWS Titan Labs, the place he builds Foundational Mannequin (FM) capabilities and options. His specialty is Pure Language Processing (NLP) and is obsessed with deep studying. Outdoors of labor he enjoys studying books and touring.

Laura Kulowski is an Utilized Scientist at Amazon’s Generative AI Innovation Heart, the place she works intently with clients to construct generative AI options. In her free time, Laura enjoys exploring new locations by bike.

Andrew Ang is a Sr. Machine Studying Engineer at AWS. Along with serving to clients construct AI/ML options, he enjoys water sports activities, squash and watching journey & meals vlogs.

Mehdi Noori is an Utilized Science Supervisor on the Generative AI Innovation Heart. With a ardour for bridging know-how and innovation, he assists AWS clients in unlocking the potential of generative AI, turning potential challenges into alternatives for fast experimentation and innovation by specializing in scalable, measurable, and impactful makes use of of superior AI applied sciences, and streamlining the trail to manufacturing.