Loading a checkpoint to make predictions

[lcoandrade]

Hi there!

I'm creating a Kaggle notebook on to use a previously saved checkpoint to make predictions.
I'm loading the checkpoint like this:

task = SemanticSegmentationTask.load_from_checkpoint('/kaggle/input/torchgeoinriacheckpoint/checkpoint.ckpt')

But when I try to make the predictions, I get this error:

RuntimeError: Input type (torch.FloatTensor) and weight type (torch.cuda.FloatTensor) should be the same or input should be a MKLDNN tensor and weight is a dense tensor

I'm creating the prediction dataset in the same way that is in my training notebook:

test_dataset = RasterDataset(
    paths='/kaggle/input/inria-aerial-image-labeling-dataset/AerialImageDataset/train/images/austin1.tif',
)
test_sampler = GridGeoSampler(test_dataset, 1024, 1024)
test_dataloader = DataLoader(test_dataset, sampler=test_sampler, collate_fn=stack_samples)

pixel_size = test_dataset.res
crs = test_dataset.crs.to_epsg()

start = timeit.default_timer() # Measuring the time
chips_generator = georreferenced_chip_generator(test_dataloader, task.model, crs, pixel_size)
print("The time taken to predict was: ", timeit.default_timer() - start)

start = timeit.default_timer() # Measuring the time
file_name = os.path.join(OUTPUT_DIR, 'merged_prediction.tif')
merge_georeferenced_chips(chips_generator, file_name)
print("The time taken to generate a georrefenced image and save it was: ", timeit.default_timer() - start)

Where georreferenced_chip_generator and merge_georeferenced_chips are functions I created to georreference predictions.

Do I need to explicitly send my input to the GPU?

[lcoandrade]

To make It work, I needed to change my function from:

def georreferenced_chip_generator(dataloader, model, crs, pixel_size):
    """
    Apply georeferencing to the predicted chip.
    
    Parameters:
        dataloader (torch.utils.data.Dataloader): Dataloader with the data to be predicted.
        model (an https://github.com/qubvel/segmentation_models.pytorch model): model used for inference.
        crs (str): Spatial Reference System (e.g., EPSG code) of the chip.
        pixel_size (float): Pixel dimensoion in map units.

    Returns:
        A list of georeferenced numpy arrays of the predicted outputs.
    """
    georref_chips_list = []
    for i, sample in enumerate(dataloader):
        image, gt_mask, bbox = sample['image'], sample['mask'], sample['bbox'][0]

        image = image/255. # as I'm not using a GeoDatamodule, I need to divide de images by 255 manually

        prediction = model.predict(image)
        prediction = torch.softmax(prediction, dim=1)
        prediction = torch.argmax(prediction, dim = 1)

        geotransform = from_origin(bbox.minx, bbox.maxy, pixel_size, pixel_size)
        
        georref_chips_list.append(create_in_memory_geochip(prediction, geotransform, crs))
    return georref_chips_list

To:

def georreferenced_chip_generator(dataloader, model, crs, pixel_size):
    """
    Apply georeferencing to the predicted chip.
    
    Parameters:
        dataloader (torch.utils.data.Dataloader): Dataloader with the data to be predicted.
        model (an https://github.com/qubvel/segmentation_models.pytorch model): model used for inference.
        crs (str): Spatial Reference System (e.g., EPSG code) of the chip.
        pixel_size (float): Pixel dimensoion in map units.

    Returns:
        A list of georeferenced numpy arrays of the predicted outputs.
    """
    georref_chips_list = []
    for i, sample in enumerate(dataloader):
        image, gt_mask, bbox = sample['image'], sample['mask'], sample['bbox'][0]

        image = image/255. # as I'm not using a GeoDatamodule, I need to divide de images by 255 manually

        prediction = model.predict(image.cuda())
        prediction = torch.softmax(prediction, dim=1)
        prediction = torch.argmax(prediction, dim = 1)

        geotransform = from_origin(bbox.minx, bbox.maxy, pixel_size, pixel_size)
        
        georref_chips_list.append(create_in_memory_geochip(prediction.cpu(), geotransform, crs))
    return georref_chips_list

My question is: Why during training I didn't need to use image.cuda() and prediction.cpu()?

[isaaccorley]

When you load the checkpoint, you need to pass map_location="cpu" like

task = SemanticSegmentationTask.load_from_checkpoint(ckpt_path, map_location="cpu")

See for more info https://pytorch.org/docs/stable/generated/torch.load.html

[lcoandrade]

Thanks!!!!!