shimmer_ssd.modules.domains.pretrained

  1from collections.abc import Mapping, Sequence
  2from pathlib import Path
  3
  4from shimmer import DomainModule, GWDecoder, GWEncoder
  5from torch.nn import Linear, Module
  6
  7from shimmer_ssd import PROJECT_DIR
  8from shimmer_ssd.ckpt_migrations import (
  9    migrate_model,
 10)
 11from shimmer_ssd.config import DomainModuleVariant, LoadedDomainConfig
 12from shimmer_ssd.errors import ConfigurationError
 13from shimmer_ssd.modules.domains.attribute import (
 14    AttributeDomainModule,
 15    AttributeLegacyDomainModule,
 16    AttributeWithUnpairedDomainModule,
 17)
 18from shimmer_ssd.modules.domains.text import GRUTextDomainModule, Text2Attr
 19from shimmer_ssd.modules.domains.visual import (
 20    VisualDomainModule,
 21    VisualLatentDomainModule,
 22    VisualLatentDomainWithUnpairedModule,
 23)
 24
 25
 26def load_pretrained_module(domain: LoadedDomainConfig) -> DomainModule:
 27    domain_checkpoint = Path(domain.checkpoint_path)
 28    module: DomainModule
 29    match domain.domain_type:
 30        case DomainModuleVariant.v:
 31            migrate_model(
 32                domain_checkpoint,
 33                PROJECT_DIR / "shimmer_ssd" / "migrations" / "visual_mod",
 34            )
 35            module = VisualDomainModule.load_from_checkpoint(
 36                domain_checkpoint, **domain.args
 37            )
 38
 39        case DomainModuleVariant.v_latents:
 40            migrate_model(
 41                domain_checkpoint,
 42                PROJECT_DIR / "shimmer_ssd" / "migrations" / "visual_mod",
 43            )
 44            v_module = VisualDomainModule.load_from_checkpoint(
 45                domain_checkpoint, **domain.args
 46            )
 47            module = VisualLatentDomainModule(v_module)
 48
 49        case DomainModuleVariant.v_latents_unpaired:
 50            migrate_model(
 51                domain_checkpoint,
 52                PROJECT_DIR / "shimmer_ssd" / "migrations" / "visual_mod",
 53            )
 54            v_module = VisualDomainModule.load_from_checkpoint(
 55                domain_checkpoint, **domain.args
 56            )
 57            module = VisualLatentDomainWithUnpairedModule(v_module)
 58
 59        case DomainModuleVariant.attr:
 60            migrate_model(
 61                domain_checkpoint,
 62                PROJECT_DIR / "shimmer_ssd" / "migrations" / "attr_mod",
 63            )
 64            module = AttributeDomainModule.load_from_checkpoint(
 65                domain_checkpoint, **domain.args
 66            )
 67
 68        case DomainModuleVariant.attr_unpaired:
 69            migrate_model(
 70                domain_checkpoint,
 71                PROJECT_DIR / "shimmer_ssd" / "migrations" / "attr_mod",
 72            )
 73            module = AttributeWithUnpairedDomainModule.load_from_checkpoint(
 74                domain_checkpoint, **domain.args
 75            )
 76
 77        case DomainModuleVariant.attr_legacy:
 78            module = AttributeLegacyDomainModule()
 79
 80        case DomainModuleVariant.t:
 81            module = GRUTextDomainModule.load_from_checkpoint(
 82                domain_checkpoint, **domain.args, strict=False
 83            )
 84            # Freezes the projector
 85            # module.embeddings.requires_grad_(False)
 86            # module.projector.requires_grad_(False)
 87
 88        case DomainModuleVariant.t_attr:
 89            assert (
 90                "text_model_path" in domain.args
 91            ), 'add "text_model_path" to the domain\'s args.'
 92            text_model = GRUTextDomainModule.load_from_checkpoint(
 93                domain.args["text_model_path"],
 94                **domain.args.get("t_args", {}),
 95            )
 96            module = Text2Attr.load_from_checkpoint(
 97                domain_checkpoint,
 98                text_model=text_model,
 99                **domain.args.get("model_args", {}),
100            )
101
102        case _:
103            raise ConfigurationError(f"Unknown domain type {domain.domain_type.name}")
104    return module
105
106
107def get_from_dict_or_val(
108    val: int | Mapping[DomainModuleVariant, int], key: DomainModuleVariant, log: str
109) -> int:
110    """
111    If val is int, return val, otherwise return val[key]
112    """
113    if isinstance(val, int):
114        return val
115
116    assert key in val, f"{key} should be defined in {log}."
117    return val[key]
118
119
120def load_pretrained_domain(
121    domain: LoadedDomainConfig,
122    workspace_dim: int,
123    encoders_hidden_dim: int | Mapping[DomainModuleVariant, int],
124    encoders_n_layers: int | Mapping[DomainModuleVariant, int],
125    decoders_hidden_dim: int | Mapping[DomainModuleVariant, int],
126    decoders_n_layers: int | Mapping[DomainModuleVariant, int],
127    is_linear: bool = False,
128    bias: bool = False,
129) -> tuple[DomainModule, Module, Module]:
130    module = load_pretrained_module(domain)
131    encoder_hidden_dim = get_from_dict_or_val(
132        encoders_hidden_dim, domain.domain_type, "global_workspace.encoders.hidden_dim"
133    )
134    decoder_hidden_dim = get_from_dict_or_val(
135        decoders_hidden_dim, domain.domain_type, "global_workspace.decoders.hidden_dim"
136    )
137
138    encoder_n_layers = get_from_dict_or_val(
139        encoders_n_layers, domain.domain_type, "global_workspace.encoder.n_layers"
140    )
141
142    decoder_n_layers = get_from_dict_or_val(
143        decoders_n_layers, domain.domain_type, "global_workspace.decoders.n_layers"
144    )
145
146    gw_encoder: Module
147    gw_decoder: Module
148    if is_linear:
149        gw_encoder = Linear(module.latent_dim, workspace_dim, bias=bias)
150        gw_decoder = Linear(workspace_dim, module.latent_dim, bias=bias)
151    else:
152        gw_encoder = GWEncoder(
153            module.latent_dim, encoder_hidden_dim, workspace_dim, encoder_n_layers
154        )
155        gw_decoder = GWDecoder(
156            workspace_dim, decoder_hidden_dim, module.latent_dim, decoder_n_layers
157        )
158
159    return module, gw_encoder, gw_decoder
160
161
162def load_pretrained_domains(
163    domains: Sequence[LoadedDomainConfig],
164    workspace_dim: int,
165    encoders_hidden_dim: int | Mapping[DomainModuleVariant, int],
166    encoders_n_layers: int | Mapping[DomainModuleVariant, int],
167    decoders_hidden_dim: int | Mapping[DomainModuleVariant, int],
168    decoders_n_layers: int | Mapping[DomainModuleVariant, int],
169    is_linear: bool = False,
170    bias: bool = False,
171) -> tuple[dict[str, DomainModule], dict[str, Module], dict[str, Module]]:
172    modules: dict[str, DomainModule] = {}
173    gw_encoders: dict[str, Module] = {}
174    gw_decoders: dict[str, Module] = {}
175    for domain in domains:
176        if domain.domain_type.kind.value.kind in modules:
177            raise ConfigurationError("Cannot load multiple domains of the same kind.")
178        model, encoder, decoder = load_pretrained_domain(
179            domain,
180            workspace_dim,
181            encoders_hidden_dim,
182            encoders_n_layers,
183            decoders_hidden_dim,
184            decoders_n_layers,
185            is_linear,
186            bias,
187        )
188        modules[domain.domain_type.kind.value.kind] = model
189        gw_encoders[domain.domain_type.kind.value.kind] = encoder
190        gw_decoders[domain.domain_type.kind.value.kind] = decoder
191    return modules, gw_encoders, gw_decoders