Skip to content

torch_world.datasets

TrajectoryDataset 

TrajectoryDataset(episodes: dict[str, Tensor], window: int)

Bases: Dataset

Sliding windows of window steps over episodes stored as tensors (N, T, ...) with a shared T.

MedMNIST 

MedMNIST(
    name: str,
    split: str = "train",
    size: int = 64,
    root: str | Path | None = None,
    download: bool = True,
)

Bases: Dataset

A MedMNIST v2 dataset: images (N, 3, size, size) uint8 and labels (N,).

The .npz file is downloaded once from Zenodo into root and checked against its MD5. Grayscale images are repeated to three channels. Items are {"image", "label"}. The data is licensed CC BY 4.0.

Yang et al., "MedMNIST v2: A large-scale lightweight benchmark for 2D and 3D biomedical image classification", Scientific Data 2023. Paper: https://arxiv.org/abs/2110.14795 Data: https://zenodo.org/records/10519652

MedicalVQA 

MedicalVQA(
    names: list[str],
    split: str = "train",
    size: int = 64,
    root: str | Path | None = None,
    download: bool = True,
    seed: int = 0,
)

Bases: Dataset

Templated visual question answering over MedMNIST datasets, in the style of the MMBU challenge.

Every image yields a modality question; datasets with more than two classes add a finding question and a verification of a named finding (true half of the time, drawn with seed); datasets with malignant classes add a malignancy question. Items are {"image", "caption", "question", "answer", "kind"}.

images holds every image once as uint8. After encode(encoder), items carry the encoder's cached outputs in place of the images, so a frozen encoder runs once per image instead of once per question.

MMBU challenge: https://akiranishii.github.io/mmbu-challenge/

info property 

info: str

A table of the source datasets and of the questions of each kind.

texts property 

texts: list[str]

Every caption, question and answer the templates can produce for these datasets.

answers property 

answers: list[str]

The candidate answers for retrieval: every answer the templates can produce.

examples 

examples(per_source: int = 4) -> tuple[Tensor, list[str]]

The first per_source images of each source dataset and their findings, for plots.images.

pixels 

pixels(items: list[int]) -> Tensor

The uint8 images of the given items.

encode 

encode(
    encoder,
    batch_size: int = 512,
    dtype=float16,
    device=None,
) -> MedicalVQA

Cache encoder(images), run on device (by default the encoder's); items then carry the outputs.

weights 

weights() -> Tensor

Per-item sampling weights that give each source dataset the same total weight.