THUDM / SwissArmyTransformer

Need more detailed information... What do you mean by "embed video encoder module from pytorch"?

Sorry! I already have a video encoder written in pytorch, how can I fully embed this module into the sat framework?

Just replace the model with your pytorch module in fine-tuning script: (Because sat models are just normal pytorch modules)

SwissArmyTransformer/examples/vit/finetune_vit_cifar10.py

Line 92 in e6abe2a

model, args = ViTFinetuneModel.from_pretrained(args.from_pretrained, args)

One more thing, maybe you need to add a disable_untrainable_params function to your model, to control what parameters you want to train:

def disable_untrainable_params(self):
    total_trainable = 0
    enable = ['mlp']
    for n, p in self.named_parameters():
        flag = False
        for e in enable:
            if e.lower() in n.lower():
                flag = True
                break
        if not flag:
            p.requires_grad_(False)
        else:
            total_trainable += p.numel()
            print_rank0(n)
    print_rank0("***** Total trainable parameters: "+str(total_trainable)+" *****")

model.disable_untrainable_params = disable_untrainable_params

How to embed video encoder module from pytorch？