Spaces:
Running
Running
File size: 4,723 Bytes
3e211d2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 | """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)
@torch.no_grad()
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()
|