The AstroCLIP model that takes an image and a spectrum and embeds them into a common space using CLIP loss. Note that you must provide the image and spectrum encoders to be used for the embedding. Args: image_encoder (nn.Module): The image encoder to be used for
(
self,
image_encoder: nn.Module,
spectrum_encoder: nn.Module,
temperature: float = 15.5,
lr: float = 1e-4,
weight_decay: float = 0.05,
epochs: int = 100,
eta_min: float = 5e-7,
logit_scale: float = 15.5,
learnable_logit_scale: bool = False,
)
| 15 | |
| 16 | class AstroClipModel(L.LightningModule): |
| 17 | def __init__( |
| 18 | self, |
| 19 | image_encoder: nn.Module, |
| 20 | spectrum_encoder: nn.Module, |
| 21 | temperature: float = 15.5, |
| 22 | lr: float = 1e-4, |
| 23 | weight_decay: float = 0.05, |
| 24 | epochs: int = 100, |
| 25 | eta_min: float = 5e-7, |
| 26 | logit_scale: float = 15.5, |
| 27 | learnable_logit_scale: bool = False, |
| 28 | ): |
| 29 | """ |
| 30 | The AstroCLIP model that takes an image and a spectrum and embeds them into a common space using CLIP loss. |
| 31 | Note that you must provide the image and spectrum encoders to be used for the embedding. |
| 32 | |
| 33 | Args: |
| 34 | image_encoder (nn.Module): The image encoder to be used for embedding. |
| 35 | spectrum_encoder (nn.Module): The spectrum encoder to be used for embedding. |
| 36 | temperature (float): The temperature parameter for the CLIP loss. |
| 37 | lr (float): The learning rate for the optimizer. |
| 38 | weight_decay (float): The weight decay for the optimizer. |
| 39 | epochs (int): The number of epochs for training. |
| 40 | eta_min (float): The minimum learning rate for the scheduler. |
| 41 | logit_scale (float): The logit scale for the CLIP loss. |
| 42 | learnable_logit_scale (bool): Whether the logit scale should be learnable. |
| 43 | """ |
| 44 | super().__init__() |
| 45 | self.save_hyperparameters() |
| 46 | |
| 47 | # Define the image and spectrum encoder |
| 48 | self.image_encoder = image_encoder |
| 49 | self.spectrum_encoder = spectrum_encoder |
| 50 | |
| 51 | # Logit scale is fixed to 15.5 and is not a learnable parameter |
| 52 | if not learnable_logit_scale: |
| 53 | self.logit_scale = np.log(logit_scale) |
| 54 | else: |
| 55 | self.logit_scale = nn.Parameter(torch.ones([]) * np.log(logit_scale)) |
| 56 | |
| 57 | # Use CLIP loss |
| 58 | self.criterion = CLIPLoss() |
| 59 | |
| 60 | def forward( |
| 61 | self, |