File size: 1,621 Bytes
8e04e6f | 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 | import torch
from models.nn.feature_factory import FeatureFactory
from .adaptive_ln_scale import AdaptiveLayerNorm
class PairReprBuilder(torch.nn.Module):
"""
Builds initial pair representation. Essentially the pair feature factory, but potentially with
an adaptive layer norm layer as well.
"""
def __init__(self, feats_repr, feats_cond, dim_feats_out, dim_cond_pair, **kwargs):
super().__init__()
self.init_repr_factory = FeatureFactory(
feats=feats_repr,
dim_feats_out=dim_feats_out,
use_ln_out=True,
mode="pair",
**kwargs,
)
self.cond_factory = None # Build a pair feature for conditioning and use it for adaln the pair representation
if feats_cond is not None:
if len(feats_cond) > 0:
self.cond_factory = FeatureFactory(
feats=feats_cond,
dim_feats_out=dim_cond_pair,
use_ln_out=True,
mode="pair",
**kwargs,
)
self.adaln = AdaptiveLayerNorm(
dim=dim_feats_out, dim_cond=dim_cond_pair
)
def forward(self, batch_nn):
mask = batch_nn["mask"] # [b, n]
pair_mask = mask[:, :, None] * mask[:, None, :] # [b, n, n]
repr = self.init_repr_factory(batch_nn) # [b, n, n, dim_feats_out]
if self.cond_factory is not None:
cond = self.cond_factory(batch_nn) # [b, n, n, dim_cond]
repr = self.adaln(repr, cond, pair_mask)
return repr
|