File size: 3,012 Bytes
396c01a | 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 | import torch
import torch.nn as nn
from torch.utils.checkpoint import checkpoint_sequential
import torchvision.models as models
def build_backbone(name: str):
if name == "efficientnet_b0":
m = models.efficientnet_b0(
weights=models.EfficientNet_B0_Weights.IMAGENET1K_V1
)
feat_dim = m.classifier[1].in_features
m.classifier = nn.Identity()
original_features = m.features
def forward_with_ckpt(x):
return checkpoint_sequential(
original_features, segments=4, input=x, use_reentrant=False
)
m.forward = lambda x: m.avgpool(forward_with_ckpt(x)).flatten(1)
return m, feat_dim
elif name == "resnet50":
m = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V1)
feat_dim = m.fc.in_features
m.fc = nn.Identity()
original_layer4 = m.layer4
def forward_resnet(x):
x = m.conv1(x); x = m.bn1(x); x = m.relu(x); x = m.maxpool(x)
x = m.layer1(x); x = m.layer2(x); x = m.layer3(x)
x = checkpoint_sequential(
original_layer4, segments=2, input=x, use_reentrant=False
)
x = m.avgpool(x)
return torch.flatten(x, 1)
m.forward = forward_resnet
return m, feat_dim
else:
return build_backbone("efficientnet_b0")
class VisualPersonalityModel(nn.Module):
def __init__(self, cfg: dict):
super().__init__()
self.backbone, feat_dim = build_backbone(cfg["backbone"])
agg_dim = feat_dim * 2
self.projector = nn.Sequential(
nn.Linear(agg_dim, cfg["embed_dim"] * 2),
nn.LayerNorm(cfg["embed_dim"] * 2),
nn.GELU(),
nn.Dropout(cfg["dropout"]),
nn.Linear(cfg["embed_dim"] * 2, cfg["embed_dim"]),
nn.LayerNorm(cfg["embed_dim"]),
)
num_traits = len(cfg["traits"])
self.regressor = nn.Sequential(
nn.Dropout(cfg["dropout"] * 0.5),
nn.Linear(cfg["embed_dim"], cfg["embed_dim"] // 2),
nn.GELU(),
nn.Dropout(cfg["dropout"] * 0.5),
nn.Linear(cfg["embed_dim"] // 2, num_traits),
nn.Sigmoid(),
)
def forward(self, frames, backbone_chunk: int = 128):
B, T, C, H, W = frames.shape
flat = frames.view(B * T, C, H, W)
feats = []
for i in range(0, flat.size(0), backbone_chunk):
feats.append(self.backbone(flat[i: i + backbone_chunk]))
x = torch.cat(feats, dim=0).view(B, T, -1)
mu = x.mean(dim=1)
std = x.std(dim=1).clamp(min=1e-6)
agg = torch.cat([mu, std], dim=-1)
emb = self.projector(agg)
pred = self.regressor(emb)
return pred, emb
def backbone_parameters(self):
return list(self.backbone.parameters())
def head_parameters(self):
return list(self.projector.parameters()) + list(self.regressor.parameters())
|