TorchScript availability

[universome]

Hi, thank you for such a nice work! I wanted to ask if by any chance you have a torchscript version of the feature extractor available? It will just make it much easier to incorporate it into existing pipelines (i would only need to replace I3D with V-JEPA and frechet distance with MMD). I tried to quickly prepare a torchscript version myself, but stumbled upon some issues:

• I noticed that the positional embeddings for e.g. vith16 do not load correctly if i use the imported V-JEPA version from your vjepa package since positional embeddings are ignored. From the shape sizes, seems like you only keep the first-frame PEs? Is it the intended behaviour?

RuntimeError: Error(s) in loading state_dict for VisionTransformer:
        size mismatch for pos_embed: copying a param with shape torch.Size([1, 1568, 1280]) from checkpoint, the shape in current model is torch.Size([1, 196, 1280]).
        size mismatch for patch_embed.proj.weight: copying a param with shape torch.Size([1280, 3, 2, 16, 16]) from checkpoint, the shape in current model is torch.Size([1280, 3, 16, 16]).

• I couldn't find the source repo for the vjepa package — is it public (from my understanding, the submitted pypi package is different from the jepa repo)?
• Judging by Fig 6, the vanilla pretrained version of V-JEPA works better, right?

[oooolga]

Hi,

• Could you please share more details about how you encountered the error? Specifically, what lines of code did you run?

• The submitted vjepa PyPI package should be identical to the previous version of the jepa package, except that it’s wrapped in the PyPI-compatible format for easier installation. (https://github.com/oooolga/v-jepa vs https://github.com/facebookresearch/jepa)

• The pretrained vjepa features are barely more aligned with human perception, based on a small-scale human study we conducted. However, due to other factors—such as sample efficiency (Figure 3 & 14), and the Metric Robustness Assessment with Progressive Distortion Level (Section 5.2)—we still recommend using the fine-tuned vjepa features for video evaluation.

[universome]

Hi, thank you for your response!

1. I am running a simple script like this as python bash-scripts/vjepa-ts.py vith16 -o snapshots/vith16.pt:

import os
import torch
from vjepa.models import vit_large, vit_huge

VJEPA_CKPT_URLS = {
    'vitl16': 'https://dl.fbaipublicfiles.com/jepa/vitl16/vitl16.pth.tar',
    'vith16': 'https://dl.fbaipublicfiles.com/jepa/vith16/vith16.pth.tar',
    'vith16-384': 'https://dl.fbaipublicfiles.com/jepa/vith16-384/vith16-384.pth.tar',
}

VIT_NAME_LOOKUP = {'vitl16': vit_large, 'vith16': vit_huge, 'vith16-384': vit_huge}

def download_ckpt(url, save_path):
    os.system(f'wget {url} -O {save_path}')

def load_model(model_name, ckpt_path, device='cuda'):
    model = VIT_NAME_LOOKUP[model_name]()
    ckpt = torch.load(ckpt_path, map_location='cpu', weights_only=False)['target_encoder']
    state_dict = {k.replace('module.', '').replace('backbone.', ''): v for k, v in ckpt.items()}
    model.load_state_dict(state_dict, strict=False)
    model.eval().to(device)
    return model

def to_torchscript(model, save_path):
    dummy_input = torch.randn(1, 3, 224, 224, device=next(model.parameters()).device)
    scripted_model = torch.jit.trace(model, dummy_input)
    scripted_model.save(save_path)

if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument('model_name', choices=VJEPA_CKPT_URLS.keys())
    parser.add_argument('--tmp_path', default='/tmp/vjepa_model.pth.tar')
    parser.add_argument('-o', '--output', required=True, help='TorchScript save path')
    args = parser.parse_args()

    if not os.path.isfile(args.tmp_path):
        download_ckpt(VJEPA_CKPT_URLS[args.model_name], args.tmp_path)
    model = load_model(args.model_name, args.tmp_path)
    to_torchscript(model, args.output)

    os.remove(args.tmp_path)
    print(f'Model {args.model_name} converted and saved to {args.output}')

The model fails at line 20 in the model.load_state_dict(...) call.

2. Oh, sorry, somehow, I missed this repo in your profile, and there is also a "Homepage" project link missing on the pypi page (e.g., like the one for this package).

3. Got it, thanks.

[oooolga]

Hi,
I'll take a look into it. In the meantime, the feature extraction code can be found here:
https://github.com/oooolga/JEDi/blob/main/tests/test_jepa_loading.py#L65

[NilanEkanayake]

I took a crack at implementing a single-page version of JEDi (requires cloned https://github.com/facebookresearch/jepa):

"""
https://github.com/oooolga/JEDi
"""
import sys
sys.path.append("jepa/")

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import os
import yaml
import requests
from tqdm import tqdm

from sklearn import metrics
from torchvision.transforms import v2, functional

from jepa.evals.video_classification_frozen.utils import ClipAggregation
from jepa.evals.video_classification_frozen.eval import init_model
from jepa.src.models.attentive_pooler import AttentiveClassifier



VJEPA_CKPT_URLS = {
    'vit_large': {
        'config': 'https://raw.githubusercontent.com/facebookresearch/jepa/refs/heads/main/configs/evals/vitl16_ssv2_16x2x3.yaml',
        'model_ckpt': 'https://dl.fbaipublicfiles.com/jepa/vitl16/vitl16.pth.tar',
        'probe_ckpt': 'https://dl.fbaipublicfiles.com/jepa/vitl16/ssv2-probe.pth.tar',
    },
}

def download(url, local_path, chunk_size=1024):
    with requests.get(url, stream=True) as r:
        total_size = int(r.headers.get("content-length", 0))
        with tqdm(total=total_size, unit="B", unit_scale=True) as pbar:
            with open(local_path, "wb") as f:
                for data in r.iter_content(chunk_size=chunk_size):
                    if data:
                        f.write(data)
                        pbar.update(chunk_size)


class JEDiMetric(nn.Module):
    def __init__(self, model='vit_large', finetuned=True, save_dir='jepa_ckpts'):
        super().__init__()
        self.finetuned = finetuned

        model_url = VJEPA_CKPT_URLS[model]['model_ckpt']
        probe_url = VJEPA_CKPT_URLS[model]['probe_ckpt']
        config_url = VJEPA_CKPT_URLS[model]['config']

        os.makedirs(save_dir, exist_ok=True)

        model_path = os.path.join(save_dir, model + '.pth.tar')
        probe_path = os.path.join(save_dir, model + '-ssv2-probe.pth.tar')
        config_path = os.path.join(save_dir, model + '.yaml')

        if not os.path.exists(model_path): # no md5 check?
            download(model_url, model_path)

        if not os.path.exists(probe_path) and finetuned:
            download(probe_url, probe_path)

        if not os.path.exists(config_path):
            download(config_url, config_path)

        with open(config_path, 'r') as y_file:
            config = yaml.load(y_file, Loader=yaml.FullLoader)

        self.model = init_model(
            "cpu",
            model_path,
            model,
            patch_size=config['pretrain']['patch_size'],
            crop_size=config['optimization']['resolution'],
            frames_per_clip=config['pretrain']['frames_per_clip'],
            tubelet_size=config['pretrain']['tubelet_size'],
            use_sdpa=config['pretrain']['use_sdpa'],
            use_SiLU=config['pretrain']['use_silu'],
            tight_SiLU=config['pretrain']['tight_silu'],
            uniform_power=config['pretrain']['uniform_power'],
        )

        self.model = ClipAggregation(self.model, tubelet_size=config['pretrain']['tubelet_size'], attend_across_segments=config['optimization']['attend_across_segments'])
        self.model.eval()

        if finetuned:
            self.classifier = AttentiveClassifier(embed_dim=self.model.embed_dim, num_heads=self.model.num_heads, depth=1, num_classes=config['data']['num_classes'])
            classifier_sd = {k.replace('module.', ''): v for k, v in torch.load(probe_path)['classifier'].items()}
            self.classifier.load_state_dict(classifier_sd)
            self.classifier.eval()
            for p in self.classifier.parameters():
                p.requires_grad = False

        for p in self.model.parameters():
            p.requires_grad = False

        self.frames_per_clip = config['pretrain']['frames_per_clip']

        self.transform = v2.Compose([
            v2.Resize(size=config['optimization']['resolution'], interpolation=functional.InterpolationMode.BICUBIC, antialias=False), # assume square already, add v2.centercrop if needed
            v2.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)), # needed?
        ])

        self.reset()
            
    def reset(self):
        self.recon_feats = []
        self.target_feats = []


    def pad_frames(self, video_tensor):
        # video_tensor: (b, c, t, h, w)
        b, c, t, h, w = video_tensor.shape
        if t >= self.frames_per_clip:
            return video_tensor
        else:
            repeated_tensor = torch.cat([video_tensor, video_tensor[:, :, -1:].repeat(1, 1, self.frames_per_clip - t, 1, 1)], dim=2)
            return repeated_tensor
    
    def update(self, recon, target): # BCTHW range -1, 1 in
        recon_feat = self.get_feats(recon)
        target_feat = self.get_feats(target)

        self.recon_feats.append(recon_feat)
        self.target_feats.append(target_feat)

    @torch.no_grad()
    def get_feats(self, videos):
        videos = (videos.clamp(-1, 1) + 1) / 2 # [-1, 1] -> [0, 1]
        videos = self.transform(videos.permute(0, 2, 1, 3, 4)).permute(0, 2, 1, 3, 4)
        videos = self.pad_frames(videos)

        feats = self.model([[videos]])[0]
        if self.finetuned: # if using ssv2 probe
            return self.classifier.pooler(feats).squeeze(1).cpu()
        else:
            return feats.mean(dim=1).cpu()

    def gather(self):
        recon_feats = torch.cat(self.recon_feats, dim=0).cpu().float().numpy()
        target_feats = torch.cat(self.target_feats, dim=0).cpu().float().numpy()
        stats = mmd_poly(target_feats, recon_feats, degree=2, coef0=0)*100 # remove x100?

        return stats
        
    def forward(self):
        pass
    
"""
FROM: https://github.com/oooolga/JEDi/blob/main/videojedi/mmd_polynomial.py
"""
def mmd_poly(X, Y, degree=2, gamma=None, coef0=0):
    """MMD using polynomial kernel (i.e., k(x,y) = (gamma <X, Y> + coef0)^degree)

    Arguments:
        X {[n_sample1, dim]} -- [X matrix]
        Y {[n_sample2, dim]} -- [Y matrix]

    Keyword Arguments:
        degree {int} -- [degree] (default: {2})
        gamma {int} -- [gamma] (default: {1})
        coef0 {int} -- [constant item] (default: {0})

    Returns:
        [scalar] -- [MMD value]
    """
    XX = metrics.pairwise.polynomial_kernel(X, X, degree, gamma, coef0)
    YY = metrics.pairwise.polynomial_kernel(Y, Y, degree, gamma, coef0)
    XY = metrics.pairwise.polynomial_kernel(X, Y, degree, gamma, coef0)
    return XX.mean() + YY.mean() - 2 * XY.mean()

@oooolga does the functionality match that of the official codebase?