<jupyter_start><jupyter_text>Image classification<jupyter_code>!pip install perceiver-io[image]==0.3.0
# Download checkpoints
!wget -nc -O logs.zip https://martin-krasser.com/perceiver/logs-update-4.zip
!unzip -qo logs.zip
import torch
import numpy as np
import matplotlib.pyplot as plt

from torchvision.datasets import MNIST
from perceiver.data.mnist import MnistPreprocessor
from perceiver.model.image.classifier import LitImageClassifier
mnist = MNIST(root='.cache', download=True, train=False)
mnist_preproc = MnistPreprocessor()
ckpt_path = 'logs/img_clf/version_0/checkpoints/epoch=014-val_loss=0.080.ckpt'

model = LitImageClassifier.load_from_checkpoint(ckpt_path).model
model.eval();
cols, rows = 3, 3
imgs = [mnist[i][0] for i in range(cols * rows)]
with torch.no_grad():
    logits = model(mnist_preproc.preprocess_batch(imgs))
    preds = logits.argmax(dim=1)
plt.figure(figsize=(8, 8))
for i, (img, pred) in enumerate(zip(imgs, preds)):
    plt.subplot(rows, cols, i + 1)
    plt.axis('off')
    plt.title(f'Prediction: {pred}')
    plt.imshow(np.array(img), cmap='gray')<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.13632286995515694, 'fraction_numerical': 0.017040358744394617, 'mean_word_length': 5.846625766871166, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '244744', 'license': 'Apache-2.0', 'n_tokens_mistral': 432, 'n_tokens_neox': 418, 'n_words': 82, 'path': 'notebooks/img-clf.ipynb', 'repo_name': 'Borda/perceiver-io', 'size': '1115'}