Runtime error on batch training. CUDA out of memory

[theoway]

Logs:

RuntimeError: CUDA out of memory. Tried to allocate 422.00 MiB (GPU 0; 11.17 GiB total capacity; 10.06 GiB already allocated; 102.88 MiB free; 10.69 GiB reserved in total by PyTorch)

Here are the P instance specs, as indicated by rastervision:

2022-09-24 18:58:07| NVIDIA-SMI 470.57.02 Driver Version: 470.57.02 CUDA Version: 11.4 |
2022-09-24 18:58:07|-------------------------------+----------------------+----------------------+
2022-09-24 18:58:042022-09-24 13:28:04:rastervision.pytorch_learner.learner: INFO - nvcc: NVIDIA (R) Cuda compiler driver
2022-09-24 18:58:04Copyright (c) 2005-2019 NVIDIA Corporation
2022-09-24 18:58:04Built on Wed_Oct_23_19:24:38_PDT_2019
2022-09-24 18:58:04Cuda compilation tools, release 10.2, V10.2.89
2022-09-24 18:58:042022-09-24 13:28:04:rastervision.pytorch_learner.learner: INFO - Python version: 3.6.13 |Anaconda, Inc.| (default, Feb 23 2021, 21:15:04)
2022-09-24 18:58:04[GCC 7.3.0]
2022-09-24 18:58:042022-09-24 13:28:04:rastervision.pytorch_learner.learner: INFO - Physical CPUs: 2
2022-09-24 18:58:042022-09-24 13:28:04:rastervision.pytorch_learner.learner: INFO - Logical CPUs: 4
2022-09-24 18:58:042022-09-24 13:28:04:rastervision.pytorch_learner.learner: INFO - Total memory: 59.96 GB
2022-09-24 18:58:042022-09-24 13:28:04:rastervision.pytorch_learner.learner: INFO - Size of /opt/data volume: 39.99 GB

Please let me know how to avoid this CUDA memory error. Here's the script of the experiment being performed:

from rastervision.core.rv_pipeline import *
from rastervision.core.backend import *
from rastervision.core.data import *
from rastervision.pytorch_backend import *
from rastervision.pytorch_learner import *

IDS = ['El_Retiro_1_3', 'El_Retiro_1_4', ....]

def build_scene(uri, id, channel_order=None):
    label_file = id + '.geojson'
    label_uri = uri + label_file
    print("Label uri: ", label_uri)

    image_file = id + '.tif'
    image_uri = uri + image_file
    print("Image uri: ", image_uri)

    raster_source = RasterioSourceConfig(
        uris=[image_uri],
        channel_order=channel_order,
        transformers=[StatsTransformerConfig()])
    
    vector_source = GeoJSONVectorSourceConfig(
        uri=label_uri,
        default_class_id=0,
        ignore_crs_field=True)

    label_source = SemanticSegmentationLabelSourceConfig(
        raster_source=RasterizedSourceConfig(
            vector_source=vector_source,
            rasterizer_config=RasterizerConfig(background_class_id=0)))
    
    label_store = SemanticSegmentationLabelStoreConfig(
        vector_output=[PolygonVectorOutputConfig(class_id=1),
                        PolygonVectorOutputConfig(class_id=2),
                        PolygonVectorOutputConfig(class_id=3)])

    return SceneConfig(
        id=id,
        raster_source=raster_source,
        label_source=label_source,
        label_store=label_store)

def get_config(runner, raw_uri, root_uri, test=True):
    scene_ids = IDS

    split_ratio = 0.8
    num_train_ids = round(len(scene_ids) * split_ratio)
    train_ids = scene_ids[0:num_train_ids]
    val_ids = scene_ids[num_train_ids:]

    num_epochs = 200

    if test:
        train_ids = train_ids[:5:]
        val_ids = val_ids[:3:]

        num_epochs = 2
    
    class_config = ClassConfig(
        names=['background', 'motorcycle', 'car', 'ghost'],
        colors=['#ffff00', '#0000ff', '#00ffff', '#00ff00']
    )

    channel_order = [0, 1, 2]
    train_scenes = [
        build_scene(raw_uri, id, channel_order) for id in train_ids
    ]
    val_scenes = [
        build_scene(raw_uri, id, channel_order) for id in val_ids
    ]

    dataset = DatasetConfig(
        class_config=class_config,
        train_scenes=train_scenes,
        validation_scenes=val_scenes)
    
    chip_sz = 650
    img_sz = chip_sz
    chip_options = SemanticSegmentationChipOptions(
        window_method=SemanticSegmentationWindowMethod.sliding,
        stride=img_sz)

    data = SemanticSegmentationImageDataConfig(
            img_sz=img_sz, num_workers=4)

    backend = PyTorchSemanticSegmentationConfig(
        data=data,
        model=SemanticSegmentationModelConfig(backbone=Backbone.resnet50),
        solver=SolverConfig(
            lr=1e-4,
            num_epochs=num_epochs,
            test_num_epochs=2,
            batch_sz=16,
            one_cycle=True))

    return SemanticSegmentationConfig(
        root_uri=root_uri,
        dataset=dataset,
        backend=backend,
        train_chip_sz=img_sz,
        predict_chip_sz=img_sz,
        chip_options=chip_options)

Should I be reducing:

• num_workers in SemanticSegmentationImageDataConfig
• batch_sz in SolverConfig
• chip_sz

Also, how many epochs should I be using? Currently, I was about to use 200, but ran into the error.

[theoway]

Also, please suggest what code should I add to ignore the null class(that is automatically created) for training.
I already have a background class. I don't want null class.

[AdeelH]

If you already have a background class, you can use it as the null class: class_config = ClassConfig(names=['background', 'foreground'], null_class='background'). Now class_config.ensure_null_class() will not add an additional null class.

[ammarsdc]

Hi @AdeelH.

I just implemented null_class in my pipeline after noticing this in the API.
May I know what will be the use case for using this?
I saw all examples in your repository does not set this even-though there is a background class.

I refer to these examples: https://github.com/azavea/raster-vision/tree/master/rastervision_pytorch_backend/rastervision/pytorch_backend/examples

[ammarsdc]

Hi @lewfish. Do you have any use cases for this? Thanks!

[lewfish]

You should explicitly set null_class if you already have a class in your labels for representing the background. Otherwise, you don't need to set it, and one will be added automatically. For more info, see https://docs.rastervision.io/en/latest/tutorials/reading_labels.html#The-%E2%80%9Cnull%E2%80%9D-class

[ammarsdc]

Great! Is this docs is newly published in the recent merged? This is really awesome. Super happy with all these docs n improvement you and your team made. Impressed and thanks!

[AdeelH]

Try reducing the img_sz (chip_sz can still be large). Reducing batch_sz can also help. You can also try to use a lighter model than the default DeepLabV3 (see isprs_potsdam.py).