wherobots/meta-tree-canopy-height-ep-torch280-cpu-pt2

First run the following to setup the environment and get the official model code

# Clone the official repo
git clone [email protected]:facebookresearch/HighResCanopyHeight.git

# Install dependencies
pip install stac-model[torch]

# Download the official pretrained checkpoints
mkdir checkpoints && aws s3 --no-sign-request sync s3://dataforgood-fb-data/forests/v1/models/saved_checkpoints/ checkpoints/

Export the model using the following:

from pathlib import Path
import sys
sys.path.append("HighResCanopyHeight")

import torch
import torch.nn as nn
import torchvision.transforms.v2 as T
from stac_model.torch.export import export, package

import src.transforms
from inference import SSLAE


# Create model and load checkpoint
class TreeCanopyHeightModel(nn.Module):
    def __init__(self, classify=True, huge=True):
        super().__init__()
        self.model = SSLAE(pretrained=None, classify=classify, huge=huge, n_bins=256)

    def forward(self, x):
        outputs = self.model(x)
        pred = 10 * outputs + 0.001
        return pred.relu()

path = "checkpoints/SSLhuge_satellite.pth"
ckpt = torch.load(path, map_location="cpu", weights_only=False)
state_dict = {f"model.{k}": v for k, v in ckpt["state_dict"].items()}
model = TreeCanopyHeightModel()
model.load_state_dict(state_dict)

# Create exportable transforms
original_transform = src.transforms.SSLNorm().Trans
norm = original_transform.transforms[-1]

transforms = nn.Sequential(
    T.Normalize(mean=[0], std=[255]),  # replace ToTensor() with normalize to 0-1
    T.Normalize(mean=norm.mean, std=norm.std)
)

# Export and save to pt2
model_program, transforms_program = export(
    input_shape=[-1, 3, 224, 224],
    model=model,
    transforms=transforms,
    device="cpu",
    dtype=torch.float32,
)
package(
    output_file=Path("model.pt2"),
    model_program=model_program,
    transforms_program=transforms_program,
    metadata_properties=None,
    aoti_compile_and_package=False
)