Rastervision (Custom loss function for torch-hub models + custom model (maybe transformer based))

[mmcs-work]

Hi,
My current training setup is exactly like the example here https://docs.rastervision.io/en/stable/usage/tutorials/train.html (using rv 0.21 in a docker image). There are two modifications I am planning to do:

1. Change the loss function in case of semantic segmentation (use diceloss along with pixel wise CE loss).
2. Change the model in future to have a good performance comparison (ViT semantic segmentation model or something similar).

• I am curious, what are the best ways of doing 1) and 2) using rastervision?

• I have looked for the second case in the issues and found out (Building model for 4-band data? #1472 (comment)) that only torch.hub models are supported. Is this the only way?

[AdeelH]

If you are using Raster Vision as a library (i.e. as shown in that tutorial), then you can use whatever model and loss function you like. You will just need to pass them to the Learner here (the SemanticSegmentationLearner constructor can also take a loss argument).

The discussion in #1472 (comment) is about using Raster Vision as a framework (see The Raster Vision Pipeline). In that scenario, the custom models or loss functions are imported via torch hub, which just means that you need to point to a GitHub repo with a hubconf.py.

[mmcs-work]

Thanks @AdeelH
You are right, I am using it as an library indeed following the example. In that case, one can create the entire model as I like since the model definition can be changed as well, along with the loss function (

raster-vision/rastervision_pytorch_learner/rastervision/pytorch_learner/semantic_segmentation_learner.py

Line 24 in 59149f8

class SemanticSegmentationLearner(Learner):

). How does this affect the model loading part later? (we train and save model bundle same way as before - and then we need to load it)

from rastervision.pytorch_learner import SemanticSegmentationLearner
class SemanticSegmentationLearnerModified(SemanticSegmentationLearner):
     def __init__():
          ....
     def train_step():...
     
     def loss():
         .... 
learner = SemanticSegmentationLearnerModified.from_model_bundle(
    model_bundle_uri='data/train-demo/model-bundle.zip',
    output_dir='data/train-demo/',
    model=model,
    train_ds=train_ds,
    valid_ds=val_ds,
    training=True,
)

If there is an example somewhere, feel free to share.

[AdeelH]

If you use a custom model or loss function, by passing it directly to the Learner constructor, you will not be able to re-construct it from the model bundle alone, as noted here:

You will need to pass it again when initializing the Learner from the model bundle:

learner = SemanticSegmentationLearner.from_model_bundle(
    model_bundle_uri='data/train-demo/model-bundle.zip',
    output_dir='data/train-demo/',
    model=my_model,
    loss=my_loss_func,
)

In the above example, the weights will be loaded from the bundle into the model that you passed in, so it's important that you pass in exactly the same model that you used for training.

---

If you do want the model to be re-constructable form the model bundle alone, you can specify a ModelConfig with an ExternalModuleConfig in LearnerConfig (instead of passing the model to the Learner directly). Example:

raster-vision/rastervision_pytorch_backend/rastervision/pytorch_backend/examples/semantic_segmentation/isprs_potsdam.py

Lines 214 to 226 in 0cac03c

	model = SemanticSegmentationModelConfig(
	external_def=ExternalModuleConfig(
	github_repo='AdeelH/pytorch-fpn:0.3',
	name='fpn',
	entrypoint='make_fpn_resnet',
	entrypoint_kwargs={
	'name': 'resnet50',
	'fpn_type': 'panoptic',
	'num_classes': num_classes,
	'fpn_channels': 256,
	'in_channels': len(channel_order),
	'out_size': (img_sz, img_sz)
	}))

You can do the same thing for the loss:

raster-vision/rastervision_pytorch_backend/rastervision/pytorch_backend/examples/chip_classification/spacenet_rio.py

Lines 163 to 171 in 0cac03c

	external_loss_def = ExternalModuleConfig(
	github_repo='AdeelH/pytorch-multi-class-focal-loss',
	name='focal_loss',
	entrypoint='focal_loss',
	force_reload=False,
	entrypoint_kwargs={
	'alpha': [.75, .25],
	'gamma': 2
	})

 solver = SolverConfig(
    ... 
    external_loss_def=external_loss_def) 

This is the same torch hub-based mechanism that is used when using RV as a framework.

[mmcs-work]

Thanks. Focusing on the first part of your answer: It should be

learner = SemanticSegmentationLearnerModified.from_model_bundle(   
    model_bundle_uri='data/train-demo/model-bundle-saved.zip',
    output_dir='data/train-demo/',
    model=my_model,
    loss=my_loss_func,
)

Since the updated class information/model definition is now in SemanticSegmentationLearnerModified which extends SemanticSegmentationLearner and the weights in model_bundle_uri ?

Edit: In your answer you said,

learner = SemanticSegmentationLearner.from_model_bundle(
    model_bundle_uri='data/train-demo/model-bundle.zip',
    output_dir='data/train-demo/',
    model=my_model,
    loss=my_loss_func,
)

In this case, my_model

class TheModelClass(nn.Module):
    def __init__(self):
        super(TheModelClass, self).__init__()
        ....

    def forward(self, x):
        ...
        return x

# Initialize model
my_model = TheModelClass()

And

class SemanticSegmentationLearnerModified(SemanticSegmentationLearner):
     def __init__():
          ....
     def train_step():...
     
     def loss():
         .... 

Because this is one of the way, we can change the loss() and how it is used etc (using SemanticSegmentationLearnerModified).

[AdeelH]

Yes, you can do all this by subclassing the SemanticSegmentationLearner. I was pointing out ways to do it without subclassing. You could also consider modifying the build_model() and the build_loss() methods. Let me know if you run into any problems with that approach.