from torch import nn, Tensor import open_clip from peft import get_peft_model, LoraConfig from ..utils import ConvRefine, ConvAdapter from ..utils import ConvUpsample, _get_norm_layer, _get_activation convnext_names_and_weights = { "convnext_base": ["laion400m_s13b_b51k"], # 107.49M "convnext_base_w": ["laion2b_s13b_b82k", "laion2b_s13b_b82k_augreg", "laion_aesthetic_s13b_b82k"], # 107.75M "convnext_base_w_320": ["laion_aesthetic_s13b_b82k", "laion_aesthetic_s13b_b82k_augreg"], # 107.75M "convnext_large_d": ["laion2b_s26b_b102k_augreg"], # 217.46M "convnext_large_d_320": ["laion2b_s29b_b131k_ft", "laion2b_s29b_b131k_ft_soup"], # 217.46M "convnext_xxlarge": ["laion2b_s34b_b82k_augreg", "laion2b_s34b_b82k_augreg_rewind", "laion2b_s34b_b82k_augreg_soup"] # 896.88M } refiner_channels = { "convnext_base": 1024, "convnext_base_w": 1024, "convnext_base_w_320": 1024, "convnext_large_d": 1536, "convnext_large_d_320": 1536, "convnext_xxlarge": 3072, } refiner_groups = { "convnext_base": 1, "convnext_base_w": 1, "convnext_base_w_320": 1, "convnext_large_d": refiner_channels["convnext_large_d"] // 512, # 3 "convnext_large_d_320": refiner_channels["convnext_large_d_320"] // 512, # 3 "convnext_xxlarge": refiner_channels["convnext_xxlarge"] // 512, # 6 } class ConvNeXt(nn.Module): def __init__( self, model_name: str, weight_name: str, block_size: int = 16, adapter: bool = False, adapter_reduction: int = 4, norm: str = "none", act: str = "none" ) -> None: super(ConvNeXt, self).__init__() assert model_name in convnext_names_and_weights, f"Model name should be one of {list(convnext_names_and_weights.keys())}, but got {model_name}." assert weight_name in convnext_names_and_weights[model_name], f"Pretrained should be one of {convnext_names_and_weights[model_name]}, but got {weight_name}." assert block_size in [32, 16, 8], f"block_size should be one of [32, 16, 8], got {block_size}" self.model_name, self.weight_name = model_name, weight_name self.block_size = block_size model = open_clip.create_model_from_pretrained(model_name, weight_name, return_transform=False).visual self.adapter = adapter if adapter: self.adapter_reduction = adapter_reduction for param in model.parameters(): param.requires_grad = False self.stem = model.trunk.stem self.depth = len(model.trunk.stages) for idx, stage in enumerate(model.trunk.stages): setattr(self, f"stage{idx}", stage) if adapter: setattr(self, f"adapter{idx}", ConvAdapter( in_channels=stage.blocks[-1].mlp.fc2.out_features, bottleneck_channels=stage.blocks[-1].mlp.fc2.out_features // adapter_reduction, ) if idx < self.depth - 1 else nn.Identity()) # No adapter for the last stage if self.model_name in ["convnext_base", "convnext_base_w", "convnext_base_w_320", "convnext_xxlarge"]: self.in_features, self.out_features = model.head.proj.in_features, model.head.proj.out_features else: # "convnext_large_d", "convnext_large_d_320": self.in_features, self.out_features = model.head.mlp.fc1.in_features, model.head.mlp.fc2.out_features if norm == "bn": norm_layer = nn.BatchNorm2d elif norm == "ln": norm_layer = nn.LayerNorm else: norm_layer = _get_norm_layer(model) if act == "relu": activation = nn.ReLU(inplace=True) elif act == "gelu": activation = nn.GELU() else: activation = _get_activation(model) if block_size == 32: self.refiner = ConvRefine( in_channels=self.in_features, out_channels=self.in_features, norm_layer=norm_layer, activation=activation, groups=refiner_groups[self.model_name], ) elif block_size == 16: self.refiner = ConvUpsample( in_channels=self.in_features, out_channels=self.in_features, norm_layer=norm_layer, activation=activation, groups=refiner_groups[self.model_name], ) else: # block_size == 8 self.refiner = nn.Sequential( ConvUpsample( in_channels=self.in_features, out_channels=self.in_features, norm_layer=norm_layer, activation=activation, groups=refiner_groups[self.model_name], ), ConvUpsample( in_channels=self.in_features, out_channels=self.in_features, norm_layer=norm_layer, activation=activation, groups=refiner_groups[self.model_name], ), ) def train(self, mode: bool = True): if self.adapter and mode: # training: self.stem.eval() for idx in range(self.depth): getattr(self, f"stage{idx}").eval() getattr(self, f"adapter{idx}").train() self.refiner.train() else: # evaluation: for module in self.children(): module.train(mode) def forward(self, x: Tensor) -> Tensor: x = self.stem(x) for idx in range(self.depth): x = getattr(self, f"stage{idx}")(x) if self.adapter: x = getattr(self, f"adapter{idx}")(x) x = self.refiner(x) return x def _convnext( model_name: str, weight_name: str, block_size: int = 16, adapter: bool = False, adapter_reduction: int = 4, lora: bool = False, lora_rank: int = 16, lora_alpha: float = 32.0, lora_dropout: float = 0.1, norm: str = "none", act: str = "none" ) -> ConvNeXt: assert not (lora and adapter), "Lora and adapter cannot be used together." model = ConvNeXt( model_name=model_name, weight_name=weight_name, block_size=block_size, adapter=adapter, adapter_reduction=adapter_reduction, norm=norm, act=act ) if lora: target_modules = [] for name, module in model.named_modules(): if isinstance(module, (nn.Linear, nn.Conv2d)) and "refiner" not in name: target_modules.append(name) lora_config = LoraConfig( r=lora_rank, lora_alpha=lora_alpha, lora_dropout=lora_dropout, bias="none", target_modules=target_modules, ) model = get_peft_model(model, lora_config) # Unfreeze refiner for name, module in model.named_modules(): if "refiner" in name: module.requires_grad_(True) return model