Better modularity for image feature test

[Necr0x0Der]

The test script https://github.com/trueagi-io/minecraft-experiments/blob/main/vla/benchmarks/block_type_classification/los_block_classifier.py
contains a lot of model-specific imports and code, and it is not always clear whether it is model-specific or general-purpose, e.g.
from extract_features import extract_dino_features_center_crop, crop_center_features, extract_full_features, base_transform
If these are model-specific functions, they should be encapsulated into the model itself, which should just take an image and produce features. Otherwise, they should be separated from the model example.

I'm not sure what TRANSFORM = transforms.Compose(...) is doing in this script. Say, normalization might not be needed by another model. And I don't see where TRANSFORM is used. It's confusing.

Besides importing all those functions, there is also a class

class Dinov2Features:
    def encode(img):
        return extract_full_features(img)

It seems that extract_full_features is only used by it, and extract_dino_features_center_crop, crop_center_features are not used. Possibly, they are supposed to replace extract_full_features in Dinov2Features for experiments.
I believe, Dinov2Features should be defined in the model class - not in the training script. All these functions should not be imported, and should be selected via Dinov2Features constructor.

I also see that base_transform (which is similar to TRANSFORM) from the model file is passed to BlockTypeDatas. It may make sense to perform image transformation at load time to save memory, and to make the preprocessing model-specific. But base_transform sounds a little confusing. Preprocessor or ImageProcessor (or just Processor) would be more conventional (e.g., transformers have AutoProcessor class for such purposes). If we suppose that the model and the processor are created in the model file to be tested, then we can agree on the following format:

from my_model import model, preprocessor

and nothing else model-specific. Then, changing my_model would be enough to test another model.

[Necr0x0Der]

Slightly different approach (which I proposed initially) is to have all the model-independent benchmarking functionality in one place, and to call it with model-specific data from another place, so it would look like this:

from los_block_classifier import benchmark
from my_model import model, preprocessor
benchmark(model, preprocessor, 'train_dataset.json', 'test_dataset.json')

In this case, there is no need for naming convention for the model and image preprocessor, and there is no need to edit a large script to benchmark another model (or on other sets). I believe that this approach is still preferable.

[DaddyWesker]

from los_block_classifier import benchmark
from my_model import model, preprocessor
benchmark(model, preprocessor, 'train_dataset.json', 'test_dataset.json')

I thought that this is exactly what is going on now:

model = Dinov2Features  # this could be replaced with any other model which has .encode(img) function

classifier = SimpleClassifier(len(classes), extract_features_size(model, train_loader)).to(device)
train_classifier(model, classifier, train_loader, test_loader)
compute_score(model, classifier, DataLoader(test_dataset, batch_size=1, shuffle=False, num_workers=NUM_WORKERS))

Here is the code from los_block_classifier.py main. You can consider main as the possible use case of the classifier. Minor difference with what you've written is that your benchmark function is the train_classifier + compute_score. If I got your proposition wrong and misimplemented it then please explain differences between current implementation and desired one and I'll change it.

I'm not sure what TRANSFORM = transforms.Compose(...) is doing in this script.

Well, I've used this transform before Oleg's code with Dino came in. Now his BaseTransform is being used. I personally have no preferences on which one should be deleted. My apologies this one was left in the code.

Besides importing all those functions, there is also a class

This class is being used as model in code in main.
model = Dinov2Features # this could be replaced with any other model which has .encode(img) function

It was used only for testing purposes. Just to show how classifier could be used. Yes, I guess this one should be placed not in the main script of classifier, but as I said main function of los_block_classifier.py is the use-case scenario.

I'll try to refactor whole code to separate files, separating classifier from model from dataset.

[DaddyWesker]

It may make sense to perform image transformation at load time to save memory

But image is being transformed currently exactly at the load time. Here, in the __getitem__ of the BlockTypeDataset:

class BlockTypeDataset(Dataset):
    def __init__(self, unpacked_dataset, num_of_classes, root_dir = "./", transform=None):
        self.root_dir = root_dir
        self.transform = transform
        self.t_set = unpacked_dataset
        self.num_of_classes = num_of_classes

    def __len__(self):
        return len(self.t_set)

    def __getitem__(self, idx):
        if torch.is_tensor(idx):
            idx = idx.tolist()
        img_name = os.path.join(self.root_dir, f"{self.t_set[idx][1]}.jpg")
        class_idx = self.t_set[idx][0]
        one_hot_vector = np.zeros(self.num_of_classes)
        one_hot_vector[class_idx] = 1
        image = io.imread(img_name)
        if self.transform is not None:
            image = self.transform(image)
        return torch.tensor(one_hot_vector), image

Or you meant something else?

[Necr0x0Der]

I thought that this is exactly what is going on now:

It's definitely not exactly what is going on now. You have a lot of code in main, it is placed together with the benchmark code, which should be model-agnostic, you have hard-coded base_transform and model assignment within main. You don't need to put all the code from main to a separate script. You should turn main into a benchmark function parameterized by the model and preprocessor.

[Necr0x0Der]

But image is being transformed currently exactly at the load time

Yes, and I say that this is ok. Another approach was to let the model to preprocess the images as it wants. Actually, if you load images and immediately put them into the model for feature extraction, you may not need to separate the preprocessor from the model.

In fact, looking into the code one more time, I'm not precisely sure that what you are doing is optimal. Why train_dataset = BlockTypeDataset(train_set_unpacked, len(classes), transform=base_transform) is doing only preprocessing without immediately calling the feature extractor? Why is feature extraction happening only here? classifier = SimpleClassifier(len(classes), extract_features_size(model, train_loader)).to(device)

[DaddyWesker]

in fact, looking into the code one more time, I'm not precisely sure that what you are doing is optimal. Why train_dataset = BlockTypeDataset(train_set_unpacked, len(classes), transform=base_transform) is doing only preprocessing without immediately calling the feature extractor?

Because this line just initializes BlockTypeDataset class to be used further in dataloader. Or your suggestion is to extract all features before train process starts and load all of them into RAM or VRAM? I'm not sure we've got so many RAM/VRAM though it should of course boost up time required for training and benchmarking.

Why is feature extraction happening only here? classifier = SimpleClassifier(len(classes), extract_features_size(model, train_loader)).to(device)

Here is not exactly feature extraction for all dataset. I'm only extracting feature for one sample of the dataset to set num_features which is used to create fully connected layers.

[Necr0x0Der]

I'm only extracting feature for one sample of the dataset to set num_features which is used to create fully connected layers.

OK, i see.