Views
No views yet

astropt directory.
pip install astropt and run the following code:1from astropt.model_utils import load_astropt
2from astropt.local_datasets import GalaxyImageDataset
3
4from datasets import load_dataset # for Smith42/galaxies
5
6import torch
7import numpy as np
8from functools import partial
9from torch.utils.data import DataLoader
10from torchvision import transforms
11
12# boilerplate to preprocess galaxy images
13def normalise(x):
14 std, mean = torch.std_mean(x, dim=1, keepdim=True)
15 return (x - mean) / (std + 1e-8)
16
17def data_transforms():
18 return transforms.Compose([transforms.Lambda(normalise)])
19
20def _process_galaxy_wrapper(idx, func):
21 """This function ensures that the image is tokenised in the same way as the pre-trained model is expecting"""
22 galaxy = func(
23 torch.from_numpy(np.array(idx["image"]).swapaxes(0, 2)).to(float)
24 ).to(torch.float)
25 galaxy_positions = torch.arange(0, len(galaxy), dtype=torch.long)
26 return {
27 "images": galaxy,
28 "images_positions": galaxy_positions,
29 }
30
31# for 095M parameter model, 015M and 850M models are also available:
32model = load_astropt("Smith42/astroPT_v2.0", path="astropt/095M")
33
34galproc = GalaxyImageDataset(
35 None,
36 spiral=True,
37 transform={"images": data_transforms()},
38 modality_registry=model.modality_registry
39)
40
41ds = (
42 load_dataset("Smith42/galaxies", split="test", revision="v2.0", streaming=True)
43 .select_columns("image")
44 .map(partial(_process_galaxy_wrapper, func=galproc.process_galaxy))
45 .with_format("torch")
46)
47
48dl = iter(DataLoader(ds, batch_size=128, num_workers=32))
49
50zs = []
51for B in dl:
52 zs.append(model.generate_embeddings(B)["images"].detach().numpy())
53zs = np.concatenate(zs)
54
55# do cool stuff with zs...