CoolFace
Modelpublic

Stemson-AI/cmmp-resnet18-512

sourceHugging Facecc-by-4.0updated 5mo agoView on Hugging Face
0likes5downloads
Model Card

CIMP: Contrastive Image-Metadata Pre-training (ResNet-18, crop 512)

A contrastive encoder that aligns HAADF-STEM microscopy images with their acquisition metadata in a shared 128-d embedding space. This variant uses a ResNet-18 image encoder trained from scratch on 512×512 patches at effective batch size 512, and is the best-performing ResNet configuration reported in the accompanying paper.

Model Details

  • —Architecture: ResNet-18 image encoder (trained from scratch, single-channel input) + 3-layer MLP metadata encoder (hidden dim 256)
  • —Embedding dimension: 128
  • —Image input: Single-channel grayscale, 512×512 pixels
  • —Metadata input: 7-d z-scored vector (pixelsize, dwelltime, convergenceangle, beamcurrent, gain, offset, innercollectionangle)
  • —Loss: Symmetric cross-entropy (CLIP-style) with learnable temperature and bias
  • —Parameters: ~11M (ResNet-18 backbone)

Retrieval Performance

Evaluated on the held-out validation split (733 images from the CMMP dataset).

MetricValue
Top-10.8594
Top-51.0000
Top-101.0000
Best epoch956 / 1000

Context among CMMP variants

VariantTop-1Top-5Top-10
This model (ResNet-18, crop 512, batch 512)0.8591.0001.000
ResNet-18, crop 2560.8440.9690.984
ViT-pretrained, crop 2560.8280.9691.000

Linear-Probe Metadata Recovery

A Ridge regression ($\alpha = 1.0$) trained on the frozen visual embedding recovers all seven acquisition parameters. Coefficient of determination ($R^2$), SMAPE (in physical units), and Pearson $r$:

Dimension$R^2$SMAPEPearson $r$
pixel_size0.74839.5%0.867
dwell_time0.81925.9%0.905
convergence_angle0.62911.6%0.793
beam_current0.69533.4%0.835
gain0.8625.0%0.929
offset0.8249.1%0.912
innercollangle0.6268.5%0.792
Mean0.74319.0%0.862

The higher SMAPE on pixelsize, dwelltime, and beam_current is expected: those dimensions are stored log10-transformed because they span several orders of magnitude in physical units, so small residuals in log-space amplify when exponentiated back.

Training Configuration

ParameterValue
DatasetCMMP HAADF-STEM (7,330 images, 6,597/733 train/val)
Image encoderResNet-18 (trained from scratch, 1-channel input)
Metadata encoder3-layer MLP, hidden dim 256
Crop size512×512 (on-the-fly from full-resolution images)
Loss functionCLIP (symmetric cross-entropy) with learnable logit bias
OptimizerAdamW (lr=1e-4, weight_decay=0.01)
SchedulerCosine annealing (LR floor ~1e-10 by epoch 1000)
Batch size64 per GPU × 8 GPUs = 512 effective
Epochs1000 (best checkpoint at epoch 956)
Hardware8× H100 GPUs

Usage

python
import torch
from models import CMMP

# Load model
model = CMMP(
    meta_input_dim=7,
    embed_dim=128,
    image_encoder="resnet18",
    image_size=512,
    meta_hidden_dim=256,
    meta_num_layers=3,
)
model.load_state_dict(torch.load("model.pth", map_location="cpu"))
model.eval()

# Embed an image and its metadata
image = torch.randn(1, 1, 512, 512)   # single-channel grayscale [0, 1]
metadata = torch.randn(1, 7)           # z-scored metadata vector

with torch.no_grad():
    img_emb, meta_emb, temp, bias = model(image, metadata)
    # img_emb: (1, 128) — L2-normalized image embedding
    # meta_emb: (1, 128) — L2-normalized metadata embedding

Files

  • —model.pth — Best checkpoint (epoch 956, highest Top-1 on val)
  • —last.pth — Final checkpoint (epoch 1000)
  • —config.json — Full training configuration (args.json from the run)
  • —training_log.csv — Per-epoch training metrics
  • —split_indices.npy — Train/val split indices (seed 67) for reproducibility
  • —linear_probe_metadata.json — Ridge-probe metadata recovery metrics

Related Models

Citation

@misc{cimp2026,
  title={Contrastive Image-Metadata Pre-training for Materials Transmission Electron Microscopy},
  author={Channing, Georgia and Keller, Debora and Rossell, Marta D. and Torr, Philip and Erni, Rolf and Helveg, Stig and Eliasson, Henrik},
  year={2026},
}