torchsample 0.1.0

Installation

Requires python >=3.8 Install torchsample via pip:

pip install torchsample

Or, if you want to install the nightly version:

pip install git+https://github.com/BrianPugh/torchsample.git@main

Training

A common scenario is to randomly sample points from a featmap and from the ground truth.

import torchsample as ts

b, c, h, w = batch["image"].shape
coords = ts.coord.rand(b, 4096, 2)  # (b, 4096, 2) where the last dim is (x, y)

featmap = feature_extractor(batch["image"])  # (b, feat, h, w)
sampled = ts.sample(coords, featmap)  # (b, 4096, feat)
gt_sample = ts.sample(coords, batch["gt"])

Inference

During inference, a comprehensive query of the network to form a complete image is common.

import torch
import torchsample as ts

b, c, h, w = batch["image"].shape
coords = ts.coord.full_like(batch["image"])
featmap = encoder(batch["image"])  # (b, feat, h, w)
feat_sampled = ts.sample(coords, featmap)  # (b, h, w, c)
output = model(featmap)  # (b, h, w, pred)
output = output.permute(0, 3, 1, 2)

Positional Encoding

Common positional encoding schemes are available.

import torchsample as ts

coords = ts.coord.rand(b, 4096, 2)
pos_enc = ts.encoding.gamma(coords)

A common task it concatenating the positional encoding to sampled values. You can do this by passing a callable into ts.sample:

import torchsample as ts

encoder = ts.encoding.Gamma()
sampled = ts.sample(coords, featmap, encoder=encoder)

Models

torchsample has some common builtin models:

import torchsample as ts

# Properly handles (..., feat) tensors.
model = ts.models.MLP(256, 256, 512, 512, 1024, 1024, 1)