Spaces:
Running
Running
| """NTv3 annotation model warm-started from NTv3-tracks, its loss and its inference.""" | |
| import functools | |
| import numpy as np | |
| import torch | |
| from torch.utils.checkpoint import checkpoint | |
| from transformers import AutoConfig, AutoModel | |
| ANNOTATION_REPO = "InstaDeepAI/NTv3_annot_plants" | |
| TRACKS_REPO = "InstaDeepAI/NTv3_tracks_1mb" | |
| # Revision of the modelling code (hosted in InstaDeepAI/NTv3_pretrained), pinned | |
| # because batch-norm running statistics depend on it. | |
| CODE_REVISION = "ea80a2df95bb605f33e1077b652da3c36c1dad56" | |
| # Which parts of the network are trained; the rest stays at the NTv3-tracks weights. | |
| TRAINABLE_PARTS = { | |
| "all": ("",), | |
| "decoder": ("model.decoder_blocks.", "model.bed_head."), | |
| "head": ("model.bed_head.",), | |
| } | |
| def set_training(model: torch.nn.Module, training: bool) -> torch.nn.Module: | |
| """Switches the trained batch norms between batch statistics and running ones. | |
| In training mode a batch norm normalises with the window's own statistics and | |
| updates its running ones. Everything else stays in eval mode: NTv3 has no | |
| dropout, and batch norms of frozen blocks keep their NTv3-tracks statistics. | |
| """ | |
| model.eval() | |
| if training and model.batch_norm == "batch": | |
| for module in model.modules(): | |
| if hasattr(module, "var_ema") and module.scale.requires_grad: | |
| module.train() | |
| return model | |
| def load_model( | |
| trainable: str = "all", | |
| gradient_checkpointing: bool = False, | |
| batch_norm: str = "frozen", | |
| ) -> torch.nn.Module: | |
| """NTv3-tracks trunk under a new 15-class, single-strand annotation head. | |
| The architecture (and head size) comes from the released plant annotation | |
| model's config; every trunk weight comes from NTv3-tracks. | |
| Args: | |
| trainable: "all", "decoder" (U-Net decoder + head) or "head". | |
| gradient_checkpointing: Recompute block activations in the backward pass. | |
| batch_norm: "frozen" keeps the NTv3-tracks running statistics; "batch" | |
| normalises trained blocks with per-window statistics during training | |
| and updates their running statistics, as our training pipeline does. | |
| """ | |
| if batch_norm == "batch" and gradient_checkpointing: | |
| # Recomputing a block in the backward pass would update its statistics twice. | |
| raise ValueError("batch_norm='batch' does not support gradient checkpointing.") | |
| config = AutoConfig.from_pretrained( | |
| ANNOTATION_REPO, trust_remote_code=True, code_revision=CODE_REVISION | |
| ) | |
| config.ntv3["keep_target_center_fraction"] = 1.0 | |
| model = AutoModel.from_pretrained( | |
| TRACKS_REPO, config=config, trust_remote_code=True, code_revision=CODE_REVISION | |
| ) | |
| prefixes = TRAINABLE_PARTS[trainable] | |
| for name, param in model.named_parameters(): | |
| param.requires_grad = name.startswith(prefixes) | |
| if gradient_checkpointing: | |
| trunk = model.model | |
| for block in [ | |
| *trunk.encoder_blocks, | |
| *trunk.mha_blocks, | |
| *trunk.mlp_blocks, | |
| *trunk.decoder_blocks, | |
| ]: | |
| block.forward = functools.partial( | |
| checkpoint, block.forward, use_reentrant=False | |
| ) | |
| model.batch_norm = batch_norm | |
| return model.eval() | |
| def logits(model: torch.nn.Module, tokens: torch.Tensor) -> torch.Tensor: | |
| """(batch, length) tokens -> (batch, length, 15) annotation logits.""" | |
| out = model(tokens, conditions_tokens=[], skip_lm_head=True) | |
| return out["bed_tracks_logits"][:, :, 0, :] | |
| def focal_loss( | |
| logits: torch.Tensor, | |
| labels: torch.Tensor, | |
| class_weights: torch.Tensor, | |
| gamma: float = 2.0, | |
| ) -> torch.Tensor: | |
| """Class-weighted focal loss averaged over labelled positions. | |
| Each position is weighted by the larger of its true and predicted class | |
| weights, so false positives of rare classes cost as much as misses. | |
| """ | |
| log_probs = torch.log_softmax(logits.float(), dim=-1) | |
| valid = labels >= 0 | |
| target = labels.clamp(min=0) | |
| log_p = log_probs.gather(-1, target[..., None])[..., 0] | |
| loss = -((1 - log_p.exp()) ** gamma) * log_p | |
| weight = torch.maximum(class_weights[target], class_weights[log_probs.argmax(-1)]) | |
| return (loss * weight * valid).sum() / valid.sum().clamp(min=1) | |
| def predict_log_probs( | |
| model: torch.nn.Module, tokens: np.ndarray, device: torch.device, dtype: torch.dtype | |
| ) -> np.ndarray: | |
| """(batch, length) tokens -> (batch, length, 15) float32 log-probabilities.""" | |
| batch = torch.from_numpy(tokens).to(device) | |
| with torch.autocast(device.type, dtype=dtype, enabled=dtype != torch.float32): | |
| out = logits(model, batch) | |
| return torch.log_softmax(out.float(), dim=-1).cpu().numpy() | |