Add hardware_dependent_fields to ZambaBaseModel's Config

[pjbull]

Currently, we publish configs with our official models so that it is easy for other people to use them. We do this by removing the properties that are hardware dependent (e.g., device: cuda or batch_size: 4) and resaving the files. This happens in a function that does the work for the configs that are passed:

zamba/zamba/models/publish_models.py

Lines 15 to 49 in 33c1c2a

	def get_model_only_params(full_configuration, subset="train_config"):
	"""Return only params that are not data or machine specific.
	Used for generating official configs.
	"""
	if subset == "train_config":
	config = full_configuration[subset]
	for key in [
	"data_dir",
	"dry_run",
	"batch_size",
	"auto_lr_find",
	"gpus",
	"num_workers",
	"max_epochs",
	"weight_download_region",
	"split_proportions",
	"save_dir",
	"overwrite",
	"skip_load_validation",
	"from_scratch",
	"model_cache_dir",
	"predict_all_zamba_species",
	]:
	config.pop(key)
	elif subset == "video_loader_config":
	config = full_configuration[subset]
	if "megadetector_lite_config" in config.keys():
	config["megadetector_lite_config"].pop("device")
	for key in ["cache_dir", "cleanup_cache"]:
	config.pop(key)
	return config

We can simplify and generalize this code by adding a hardware_dependent_fields property to the config and then a hardware_independent_dict method to the model that will serialize the config, removing all the hardware_dependent_fields on that model and any children that inherit from the same base model. Here's a working implementation with a small example config:

from pydantic import BaseModel


class CustomBase(BaseModel):
    class Config:
        hardware_dependent = []
        
    def dict_hardware_independent(self):
        full_dict = self.dict()
        
        # remove hardware dependent fields on children
        for field_name, field_type in self.__fields__.items():
            field_value = getattr(self, field_name, None)
            if isinstance(field_value, CustomBase):
                full_dict[field_name] = field_value.dict_hardware_independent()

        # remove hardware dependent fields on this model
        for f in self.Config.hardware_dependent:
            full_dict.pop(f)
                
        return full_dict
        

class SubModelC(CustomBase):
    x: str

    class Config:
        hardware_dependent = ['x']
        
class SubModelA(CustomBase):
    a: str
    b: str
    c: SubModelC

    class Config:
        hardware_dependent = ['a']
        
instance_a = SubModelA(a="a", b="b", c=SubModelC(x='x'))

instance_a.dict_hardware_independent()
# > {'b': 'b', 'c': {}}