Views
No views yet
1import torch
2from pathlib import Path
3
4# Load model checkpoint
5model_path = "astropt/090M/ckpt.pt"
6device = "cuda" if torch.cuda.is_available() else "cpu"
7
8# Load state dict
9checkpoint = torch.load(model_path, map_location=device)
10# Initialize your model architecture here
11# model.load_state_dict(checkpoint)1from datasets import load_dataset
2import torch
3import numpy as np
4
5# Load dataset
6dataset = load_dataset(
7 "msiudek/astroPT_euclid_dataset",
8 split="train_batch_1",
9 streaming=True
10)
11
12# Run inference
13model.eval()
14with torch.no_grad():
15 for sample in dataset:
16 # Stack 4 imaging bands: VIS + Y + J + H
17 images = torch.stack([
18 torch.tensor(sample['VIS_image'], dtype=torch.float32),
19 torch.tensor(sample['NISP_Y_image'], dtype=torch.float32),
20 torch.tensor(sample['NISP_J_image'], dtype=torch.float32),
21 torch.tensor(sample['NISP_H_image'], dtype=torch.float32),
22 ]) # [4, 224, 224]
23 images = images.unsqueeze(0) # [1, 4, 224, 224]
24
25 # 13-band SED
26 sed = torch.tensor(sample['SED'], dtype=torch.float32)
27 sed = sed.unsqueeze(0) # [1, 13]
28
29 # Handle NaN values in SED
30 sed = torch.nan_to_num(sed, nan=0.0)
31
32 # Get multi-modal embeddings
33 embeddings = model(images, sed)1@article{Siudek2025,
2 title={AstroPT: Astronomical Physics Transformers for Multi-modal Learning},
3 author={Siudek, M et al},
4 journal={Euclid Collaboration},
5 eprint={2503.15312},
6 archivePrefix={arXiv},
7 year={2025},
8 url={https://ui.adsabs.harvard.edu/abs/2025arXiv250315312E/abstract}
9}