How to load the model as a PyTorch module?

[yuxi-liu-wired]

I am trying to load just the downloaded model checkpoint.pth to be used as a PyTorch module in another model. I found it difficult to extract the minimum amount of code necessary to load checkpoint.pth as a PyTorch module.

Specifically, I cannot just run

python main_sim.py --dataset wikiart -a vit_large --pt_style csd --feattype normal --world-size 1 
--dist-url tcp://localhost:6001 -b 128 -j 8 --embed_dir ./embeddings --data-dir <path to WikiArt dataset>
--model_path <path to CSD weights>

because I need to use the module as part of a larger model, not on its own.

Further, there are two discrepancies:

• It is claimed that the checkpoint.pth available for download is a ViT-L, but it has only 304M parameters, while ViT-L has 428M parameters.
• The paper says you have finetuned both ViT-B and ViT-L, but there is no ViT-B checkpoint available for download.

[learn2phoenix]

@yuxi-liu-wired Thanks for showing interest in our work. We've some commitments and intend to respond to you towards the first week of October. I apologize if this impacts your workflow.

[yuxi-liu-wired]

I have figured out some things.

1. The CSD in the title seems to also mean "content-style disentanglement". This can be seen by the use of gradient reversal and alpha.
2. The 305M is just the ViT-L/14 image encoder half of OpenAI CLIP ViT-L/14, without the text encoder.
3. The model has TWO heads: a style head and a content head. Each is a $1024 \times 768$ projection matrix. It outputs two $768$-dimensional vector.
4. The model is constructed by removing the ViT-L/14 image encoder's projection head, then appending the two new heads.
5. The model can be loaded as follows.

I have converted the model into several formats, including safetensor, onnx, and pytorch without the optimizer states. It is for download in https://huggingface.co/yuxi-liu-wired/CSD/tree/main. For example, downloading csd_clip_model_pytorch.pt and putting it in ../model_weight/csd_clip_model_pytorch.pt would allow the following script to work.

import copy
import os
import io
import torch
import torch.nn as nn
import clip
import pandas as pd
from PIL import Image
from tqdm import tqdm
import numpy as np
from transformers import CLIPProcessor, CLIPVisionModel


class CSD_CLIP(nn.Module):
    """backbone + projection head"""
    def __init__(self, name='vit_large',content_proj_head='default'):
        super(CSD_CLIP, self).__init__()
        self.content_proj_head = content_proj_head
        if name == 'vit_large':
            clipmodel, _ = clip.load("ViT-L/14")
            self.backbone = clipmodel.visual
            self.embedding_dim = 1024
        elif name == 'vit_base':
            clipmodel, _ = clip.load("ViT-B/16")
            self.backbone = clipmodel.visual
            self.embedding_dim = 768 
            self.feat_dim = 512
        else:
            raise Exception('This model is not implemented')

        self.last_layer_style = copy.deepcopy(self.backbone.proj)
        self.last_layer_content = copy.deepcopy(self.backbone.proj)

        self.backbone.proj = None

    @property
    def dtype(self):
        return self.backbone.conv1.weight.dtype
    
    @property
    def device(self):
        return next(self.parameters()).device

    def forward(self, input_data):
        
        feature = self.backbone(input_data)

        style_output = feature @ self.last_layer_style
        style_output = nn.functional.normalize(style_output, dim=1, p=2)

        content_output = feature @ self.last_layer_content
        content_output = nn.functional.normalize(content_output, dim=1, p=2)
        
        return feature, content_output, style_output


def get_model_processor(name):
    if name == "CLIP":
        model = CLIPVisionModel.from_pretrained("openai/clip-vit-large-patch14")
        processor = CLIPProcessor.from_pretrained("openai/clip-vit-large-patch14")
        model = model.to(device)
    elif name == "CSD":
        model = CSD_CLIP("vit_large")
        model.load_state_dict(torch.load('../model_weight/csd_clip_model_pytorch.pt', 
                                         map_location=device, 
                                         weights_only=True))
        model.eval()
        model = model.to(device)
        processor = CLIPProcessor.from_pretrained("openai/clip-vit-large-patch14")
    return model, processor


def process_batch(image_paths, model, processor, model_name):
    images = [Image.open(img_path).convert("RGB") for img_path in image_paths]
    inputs = processor(images=images, return_tensors="pt", padding=True, truncation=True)
    inputs = {k: v.to(model.dtype) for k, v in inputs.items()}

    if model_name == "CLIP":
        with torch.no_grad():
            outputs = model(inputs['pixel_values'].to(model.device))
        
        return outputs.last_hidden_state[:, 0, :].cpu().numpy()
    elif model_name == "CSD":
        with torch.no_grad():
            _, content_output, style_output = model(inputs['pixel_values'].to(model.device))
        
        return content_output.cpu().numpy(), style_output.cpu().numpy()
    else:
        raise Exception(f'Model {model_name} not implemented')


def to_jpeg(image):
    buffered = io.BytesIO()
    image.save(buffered, format='JPEG')
    return buffered.getvalue() 

I have tried running the model on my own dataset, then using a t-SNE embedding. It seems to me that both the style head and the content head does a similarly good job. Both are notably better than CLIP.

[learn2phoenix]

@yuxi-liu-wired Many thanks for your efforts, and for sharing this analysis. We will try to merge these changes soon. If you want, please go ahead and raise a PR

[yuxi-liu-wired]

I had uploaded the model on HuggingFace https://huggingface.co/yuxi-liu-wired/CSD along with the code usage examples I gave above. Please take a look.