LLoCa-ParticleTransformer ========================= We start with the updated ParticleTransformer implementation, available at https://github.com/hqucms/weaver-core/blob/dev/custom_train_eval/weaver/nn/model/ParticleTransformer.py, and made small changes to improve numerical stability and efficiency. Similar to :doc:`transformer`, we use :class:`~lloca.backbone.attention.LLoCaAttention` to efficiently implement tensorial message-passing. ParticleTransformer features two types of attention blocks, particle self-attention as the main workhorse and class attention for the final aggregation over the jet. We find that it is sufficient to use tensorial message-passing only in the particle self-attention blocks. .. code-block:: diff import copy import math import random from collections.abc import Callable from functools import partial from typing import Any, Optional, Tuple import torch import torch.nn as nn import torch.nn.functional as F +from ..reps.tensorreps import TensorReps +from .attention import LLoCaAttention @torch.jit.script def delta_phi(a, b): return (a - b + math.pi) % (2 * math.pi) - math.pi @torch.jit.script def delta_r2(eta1, phi1, eta2, phi2): return (eta1 - eta2) ** 2 + delta_phi(phi1, phi2) ** 2 def to_pt2(x, eps=1e-8): pt2 = x[:, :2].square().sum(dim=1, keepdim=True) if eps is not None: pt2 = pt2.clamp(min=eps) return pt2 def to_m2(x, eps=1e-8): m2 = x[:, 3:4].square() - x[:, :3].square().sum(dim=1, keepdim=True) if eps is not None: m2 = m2.clamp(min=eps) return m2 def to_ptrapphim(x, return_mass=True, eps=1e-8): # x: (N, 4, ...), dim1 : (px, py, pz, E) px, py, pz, energy = x.split((1, 1, 1, 1), dim=1) pt = torch.sqrt(to_pt2(x, eps=eps)) # rapidity = 0.5 * torch.log((energy + pz) / (energy - pz)) rapidity = 0.5 * torch.log((1 + (2 * pz) / (energy - pz).clamp(min=1e-20)).clamp(min=1e-20)) phi = torch.atan2(py, px) if not return_mass: return torch.cat((pt, rapidity, phi), dim=1) else: m = torch.sqrt(to_m2(x, eps=eps)) return torch.cat((pt, rapidity, phi, m), dim=1) def boost(x, boostp4, eps=1e-8): # boost x to the rest frame of boostp4 # x: (N, 4, ...), dim1 : (px, py, pz, E) p3 = -boostp4[:, :3] / boostp4[:, 3:].clamp(min=eps) b2 = p3.square().sum(dim=1, keepdim=True) gamma = (1 - b2).clamp(min=eps) ** (-0.5) gamma2 = (gamma - 1) / b2 gamma2.masked_fill_(b2 == 0, 0) bp = (x[:, :3] * p3).sum(dim=1, keepdim=True) v = x[:, :3] + gamma2 * bp * p3 + x[:, 3:] * gamma * p3 return v def p3_norm(p, eps=1e-8): return p[:, :3] / p[:, :3].norm(dim=1, keepdim=True).clamp(min=eps) def to_energy_momentum(x, return_unit_vector=True): energy = x[:, 3:4] p2 = x[:, :3].square().sum(dim=1, keepdim=True) mom = torch.sqrt(p2) if return_unit_vector: return energy, mom, x[:, :3] / mom.clamp(min=1e-8) else: return energy, mom def to_cos_sin_angles(xi, xj, normed_inputs=False, eps=1e-8): if normed_inputs: ni, nj = xi, xj else: ni, nj = p3_norm(xi, eps), p3_norm(xj, eps) cos = (ni * nj).sum(dim=1, keepdim=True).clamp(min=-1, max=1) sin = torch.linalg.cross(ni, nj, dim=1).norm(dim=1, keepdim=True).clamp(min=0, max=1) return cos, sin def pairwise_lv_fts_pp(xi, xj, num_outputs=4, eps=1e-8): pti, rapi, phii = to_ptrapphim(xi, False, eps=None).split((1, 1, 1), dim=1) ptj, rapj, phij = to_ptrapphim(xj, False, eps=None).split((1, 1, 1), dim=1) # modified this for convenience (only lorentz scalars is most conservative) xij = xi + xj lnm2 = torch.log(to_m2(xij, eps=eps)) if num_outputs > 0: outputs = [lnm2] if num_outputs > 1: delta = delta_r2(rapi, phii, rapj, phij).sqrt() lndelta = torch.log(delta.clamp(min=eps)) ptmin = torch.minimum(pti, ptj) lnkt = torch.log((ptmin * delta).clamp(min=eps)) lnz = torch.log((ptmin / (pti + ptj).clamp(min=eps)).clamp(min=eps)) outputs += [lnkt, lnz, lndelta] if num_outputs > 4: lnds2 = torch.log(torch.clamp(-to_m2(xi - xj, eps=None), min=eps)) outputs.append(lnds2) # the following features are not symmetric for (i, j) if num_outputs > 5: xj_boost = boost(xj, xij) costheta = (p3_norm(xj_boost, eps=eps) * p3_norm(xij, eps=eps)).sum(dim=1, keepdim=True) outputs.append(costheta) if num_outputs > 6: deltarap = rapi - rapj deltaphi = delta_phi(phii, phij) outputs += [deltarap, deltaphi] assert len(outputs) == num_outputs return torch.cat(outputs, dim=1) def pairwise_lv_fts_ee(xi, xj, num_outputs=6, eps=1e-8): # outputs: [lnm2, cos_angle, sin_angle, lnkt, lnz, lnjade] lnm2 = torch.log(to_m2(xi + xj, eps=eps)) outputs = [lnm2] if num_outputs > 1: ei, pi, ni = to_energy_momentum(xi) ej, pj, nj = to_energy_momentum(xj) cos_angle, sin_angle = to_cos_sin_angles(ni, nj, normed_inputs=True) outputs += [cos_angle, sin_angle] if num_outputs > 3: pmin = torch.minimum(pi, pj) lnkt = torch.log((pmin * sin_angle).clamp(min=eps)) lnz = torch.log((pmin / (pi + pj).clamp(min=eps)).clamp(min=eps)) outputs += [lnkt, lnz] if num_outputs > 5: lnjade = torch.log((ei * ej * (1 - cos_angle)).clamp(min=eps)) outputs.append(lnjade) assert len(outputs) == num_outputs return torch.cat(outputs, dim=1) def build_sparse_tensor(uu, idx, seq_len): # inputs: uu (N, C, num_pairs), idx (N, 2, num_pairs) # return: (N, C, seq_len, seq_len) batch_size, num_fts, num_pairs = uu.size() idx = torch.min(idx, torch.ones_like(idx) * seq_len) i = torch.cat( ( torch.arange(0, batch_size, device=uu.device) .repeat_interleave(num_fts * num_pairs) .unsqueeze(0), torch.arange(0, num_fts, device=uu.device) .repeat_interleave(num_pairs) .repeat(batch_size) .unsqueeze(0), idx[:, :1, :].expand_as(uu).flatten().unsqueeze(0), idx[:, 1:, :].expand_as(uu).flatten().unsqueeze(0), ), dim=0, ) return torch.sparse_coo_tensor( i, uu.flatten(), size=(batch_size, num_fts, seq_len + 1, seq_len + 1), device=uu.device, ).to_dense()[:, :, :seq_len, :seq_len] def tril_indices(row, col, offset=0, *, dtype=torch.long, device="cpu"): return torch.ones(row, col, dtype=dtype, device=device).tril(offset).nonzero().T class SequenceTrimmer(nn.Module): def __init__(self, enabled=False, target=(0.9, 1.02), warmup_steps=5, **kwargs) -> None: super().__init__(**kwargs) self.enabled = enabled self.target = target self.warmup_steps = warmup_steps self.register_buffer("_counter", torch.LongTensor([0]), persistent=False) def forward(self, x, v=None, mask=None, uu=None): # x: (N, C, P) # v: (N, 4, P) [px,py,pz,energy] # mask: (N, 1, P) -- real particle = 1, padded = 0 # uu: (N, C', P, P) if mask is None: mask = torch.ones_like(x[:, :1]) mask = mask.bool() if self.enabled: if self._counter < self.warmup_steps: self._counter.add_(1) else: if v is not None: if not isinstance(v, (list, tuple)): v = [v] if self.training: q = min(1, random.uniform(*self.target)) maxlen = torch.quantile(mask.float().sum(dim=-1), q).long() rand = torch.rand_like(mask.float()) rand.masked_fill_(~mask, -1) perm = rand.argsort(dim=-1, descending=True) # (N, 1, P) mask = torch.gather(mask, -1, perm) x = torch.gather(x, -1, perm.expand_as(x)) if v is not None: v = [torch.gather(_v, -1, perm.expand_as(_v)) for _v in v] if uu is not None: uu = torch.gather(uu, -2, perm.unsqueeze(-1).expand_as(uu)) uu = torch.gather(uu, -1, perm.unsqueeze(-2).expand_as(uu)) else: maxlen = mask.sum(dim=-1).max() maxlen = max(maxlen, 1) if maxlen < mask.size(-1): mask = mask[:, :, :maxlen] x = x[:, :, :maxlen] if v is not None: v = [_v[:, :, :maxlen] for _v in v] if uu is not None: uu = uu[:, :, :maxlen, :maxlen] if v is not None: if len(v) == 1: v = v[0] return x, v, mask, uu class SwiGLUFFN(nn.Module): def __init__( self, in_features: int, hidden_features: int | None = None, out_features: int | None = None, drop: float = 0.0, bias: bool = True, ) -> None: super().__init__() hidden_features = hidden_features or in_features out_features = out_features or in_features self.w12 = nn.Linear(in_features, 2 * hidden_features, bias=bias) self.drop = nn.Dropout(drop) self.w3 = nn.Linear(hidden_features, out_features, bias=bias) def forward(self, x: torch.Tensor) -> torch.Tensor: x12 = self.w12(x) x1, x2 = x12.chunk(2, dim=-1) hidden = F.silu(x1) * x2 hidden = self.drop(hidden) return self.w3(hidden) class Embed(nn.Module): def __init__(self, input_dim, dims, normalize_input=True, activation="gelu"): super().__init__() self.input_bn = nn.BatchNorm1d(input_dim) if normalize_input else None module_list = [] for dim in dims: module_list.extend( [ nn.LayerNorm(input_dim), nn.Linear(input_dim, dim), nn.GELU() if activation == "gelu" else nn.ReLU(), ] ) input_dim = dim self.embed = nn.Sequential(*module_list) def forward(self, x): if self.input_bn is not None: # x: (batch, embed_dim, seq_len) x = self.input_bn(x) x = x.transpose(1, 2).contiguous() # x: (batch, seq_len, embed_dim) return self.embed(x) class PairEmbed(nn.Module): def __init__( self, pairwise_lv_dim, pairwise_input_dim, dims, pairwise_lv_type="pp", remove_self_pair=False, use_pre_activation_pair=True, normalize_input=True, activation="gelu", eps=1e-8, for_onnx=False, sparse_eval=None, ): super().__init__() self.pairwise_lv_dim = pairwise_lv_dim self.pairwise_input_dim = pairwise_input_dim self.remove_self_pair = remove_self_pair self.for_onnx = for_onnx self.sparse_eval = (not for_onnx) if sparse_eval is None else sparse_eval self.out_dim = dims[-1] if pairwise_lv_type == "pp": self.is_symmetric = (pairwise_lv_dim <= 5) and (pairwise_input_dim == 0) self.pairwise_lv_fts = partial(pairwise_lv_fts_pp, num_outputs=pairwise_lv_dim, eps=eps) elif pairwise_lv_type == "ee": self.is_symmetric = (pairwise_lv_dim <= 6) and (pairwise_input_dim == 0) self.pairwise_lv_fts = partial(pairwise_lv_fts_ee, num_outputs=pairwise_lv_dim, eps=eps) else: raise RuntimeError("Invalid value for `pairwise_lv_type`: " + pairwise_lv_type) if pairwise_lv_dim > 0: input_dim = pairwise_lv_dim module_list = [nn.BatchNorm1d(input_dim)] if normalize_input else [] for dim in dims: module_list.extend( [ nn.Conv1d(input_dim, dim, 1), nn.BatchNorm1d(dim), nn.GELU() if activation == "gelu" else nn.ReLU(), ] ) input_dim = dim if use_pre_activation_pair: module_list = module_list[:-1] self.embed = nn.Sequential(*module_list) if pairwise_input_dim > 0: input_dim = pairwise_input_dim module_list = [nn.BatchNorm1d(input_dim)] if normalize_input else [] for dim in dims: module_list.extend( [ nn.Conv1d(input_dim, dim, 1), nn.BatchNorm1d(dim), nn.GELU() if activation == "gelu" else nn.ReLU(), ] ) input_dim = dim if use_pre_activation_pair: module_list = module_list[:-1] self.fts_embed = nn.Sequential(*module_list) def _forward_dense(self, x, uu=None, mask=None): # x: (batch, v_dim, seq_len) # uu: (batch, v_dim, seq_len, seq_len) assert x is not None or uu is not None with torch.no_grad(): if x is not None: batch_size, _, seq_len = x.size() else: batch_size, _, seq_len, _ = uu.size() if self.is_symmetric: tril_indices_fn = tril_indices if self.for_onnx else torch.tril_indices i, j = tril_indices_fn( seq_len, seq_len, offset=-1 if self.remove_self_pair else 0, device=(x if x is not None else uu).device, ) if x is not None: x = x.unsqueeze(-1).repeat(1, 1, 1, seq_len) xi = x[:, :, i, j] # (batch, dim, seq_len*(seq_len+1)/2) xj = x[:, :, j, i] x = self.pairwise_lv_fts(xi, xj) if uu is not None: # (batch, dim, seq_len*(seq_len+1)/2) uu = uu[:, :, i, j] else: if x is not None: x = self.pairwise_lv_fts(x.unsqueeze(-1), x.unsqueeze(-2)) if self.remove_self_pair: i = torch.arange(0, seq_len, device=x.device) x[:, :, i, i] = 0 x = x.view(-1, self.pairwise_lv_dim, seq_len * seq_len) if uu is not None: uu = uu.view(-1, self.pairwise_input_dim, seq_len * seq_len) # with grad elements = 0 if x is not None: elements = elements + self.embed(x) if uu is not None: elements = elements + self.fts_embed(uu) if self.is_symmetric: y = torch.zeros( batch_size, self.out_dim, seq_len, seq_len, dtype=elements.dtype, device=elements.device, ) y[:, :, i, j] = elements y[:, :, j, i] = elements else: y = elements.view(-1, self.out_dim, seq_len, seq_len) return y def _forward_sparse(self, x, uu=None, mask=None): # x: (batch, v_dim, seq_len) # uu: (batch, v_dim, seq_len, seq_len) assert x is not None or uu is not None with torch.no_grad(): if x is not None: batch_size, _, seq_len = x.size() else: batch_size, _, seq_len, _ = uu.size() i0, i1, i2, i3 = (Ellipsis,) * 4 if mask is not None: mask = mask.unsqueeze(-1) * mask.unsqueeze(-2) # (batch_size, 1, seq_len, seq_len) if self.is_symmetric: offset = -1 if self.remove_self_pair else 0 i0, _, i2, i3 = mask.float().tril(offset).nonzero(as_tuple=True) else: i0, _, i2, i3 = mask.nonzero(as_tuple=True) if x is not None: x = self.pairwise_lv_fts(x.unsqueeze(-1), x.unsqueeze(-2)) x = x.permute(0, 2, 3, 1)[i0, i2, i3, :] # (num_elements, pairwise_lv_dim) x = x.T.unsqueeze(0).contiguous() # (1, pairwise_lv_dim, num_elements) if uu is not None: uu = uu.permute(0, 2, 3, 1)[i0, i2, i3, :] # (num_elements, pairwise_input_dim) uu = uu.T.unsqueeze(0).contiguous() # (1, pairwise_input_dim, num_elements) # with grad elements = 0 if x is not None: elements = elements + self.embed(x) if uu is not None: elements = elements + self.fts_embed(uu) elements = elements.squeeze(0).T # (num_elements, out_dim) y = torch.zeros( batch_size, seq_len, seq_len, self.out_dim, dtype=elements.dtype, device=elements.device, ) y[i0, i2, i3, :] = elements if self.is_symmetric: y[i0, i3, i2, :] = elements y = y.permute(0, 3, 1, 2).contiguous() return y def forward(self, x, uu=None, mask=None): if self.sparse_eval: return self._forward_sparse(x, uu=uu, mask=mask) else: return self._forward_dense(x, uu=uu, mask=mask) def _canonical_mask( mask: torch.Tensor | None, mask_name: str, other_type: Any | None, other_name: str, target_type: Any, check_other: bool = True, ) -> torch.Tensor | None: if mask is not None: _mask_dtype = mask.dtype _mask_is_float = torch.is_floating_point(mask) if _mask_dtype != torch.bool and not _mask_is_float: raise AssertionError(f"only bool and floating types of {mask_name} are supported") if not _mask_is_float: mask = torch.zeros_like(mask, dtype=target_type).masked_fill_(mask, float("-inf")) return mask def _none_or_dtype(input: torch.Tensor | None): if input is None: return None elif isinstance(input, torch.Tensor): return input.dtype raise RuntimeError("input to _none_or_dtype() must be None or torch.Tensor") class Attention(torch.nn.Module): def __init__( self, + attention, embed_dim, num_heads, dropout=0.0, bias=True, device=None, dtype=None, ) -> None: factory_kwargs = {"device": device, "dtype": dtype} super().__init__() self.embed_dim = embed_dim self.num_heads = num_heads self.dropout = dropout self.head_dim = embed_dim // num_heads assert self.head_dim * num_heads == self.embed_dim, ( "embed_dim must be divisible by num_heads" ) self.in_proj = torch.nn.Linear(embed_dim, 3 * embed_dim, bias=bias, **factory_kwargs) self.out_proj = torch.nn.Linear(embed_dim, embed_dim, bias=bias, **factory_kwargs) + self.attention = attention def _load_from_state_dict( self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs, ): for k in state_dict.keys(): if k.endswith("in_proj_weight"): state_dict[k.replace("_weight", ".weight")] = state_dict.pop(k) elif k.endswith("in_proj_bias"): state_dict[k.replace("_bias", ".bias")] = state_dict.pop(k) super()._load_from_state_dict( state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs, ) def forward( self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, key_padding_mask: torch.Tensor | None = None, attn_mask: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor | None]: bsz, tgt_len, _ = query.shape _, src_len, _ = key.shape # (bsz, src_len) key_padding_mask = _canonical_mask( mask=key_padding_mask, mask_name="key_padding_mask", other_type=_none_or_dtype(attn_mask), other_name="attn_mask", target_type=query.dtype, ) # (bsz, num_heads, tgt_len, src_len) attn_mask = _canonical_mask( mask=attn_mask, mask_name="attn_mask", other_type=None, other_name="", target_type=query.dtype, check_other=False, ) # merge key padding and attention masks if key_padding_mask is not None: assert key_padding_mask.shape == ( bsz, src_len, ), ( f"expecting key_padding_mask shape of {(bsz, src_len)}, but got {key_padding_mask.shape}" ) key_padding_mask = key_padding_mask.view(bsz, 1, 1, src_len).expand( -1, self.num_heads, -1, -1 ) if attn_mask is None: attn_mask = key_padding_mask else: assert attn_mask.shape == ( bsz, self.num_heads, tgt_len, src_len, ), ( f"expecting attn_mask shape of {(bsz, self.num_heads, tgt_len, src_len)}, but got {attn_mask.shape}" ) attn_mask = attn_mask + key_padding_mask # (bsz, seq_len, num_heads*head_dim) q, k, v = F._in_projection_packed(query, key, value, self.in_proj.weight, self.in_proj.bias) # -> (bsz, num_heads, src/tgt_len, head_dim) q = q.view(bsz, tgt_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous() k = k.view(bsz, src_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous() v = v.view(bsz, src_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous() dropout_p = self.dropout if self.training else 0.0 - attn_output = F.scaled_dot_product_attention(q, k, v, attn_mask, dropout_p) + if self.attention is not None: + # particle attention + attn_output = self.attention( + q, + k, + v, + attn_mask=attn_mask, + dropout_p=dropout_p, + ) + else: + # class token attention + attn_output = F.scaled_dot_product_attention(q, k, v, attn_mask, dropout_p) attn_output = attn_output.transpose(1, 2).reshape(bsz, tgt_len, self.embed_dim) attn_output = self.out_proj(attn_output) return attn_output, None class LayerScale(nn.Module): def __init__( self, dim: int, init_values: float = 1e-5, inplace: bool = False, ) -> None: super().__init__() self.inplace = inplace self.gamma = nn.Parameter(init_values * torch.ones(dim)) def forward(self, x: torch.Tensor) -> torch.Tensor: return x.mul_(self.gamma) if self.inplace else x * self.gamma def drop_path(x, drop_prob: float = 0.0, training: bool = False, scale_by_keep: bool = True): """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). This is the same as the DropConnect impl I created for EfficientNet, etc networks, however, the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper... See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use 'survival rate' as the argument. """ if drop_prob == 0.0 or not training: return x keep_prob = 1 - drop_prob shape = (x.shape[0],) + (1,) * (x.ndim - 1) # work with diff dim tensors, not just 2D ConvNets random_tensor = x.new_empty(shape).bernoulli_(keep_prob) if keep_prob > 0.0 and scale_by_keep: random_tensor.div_(keep_prob) return x * random_tensor class DropPath(nn.Module): """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).""" def __init__(self, drop_prob: float = 0.0, scale_by_keep: bool = True): super(DropPath, self).__init__() self.drop_prob = drop_prob self.scale_by_keep = scale_by_keep def forward(self, x): return drop_path(x, self.drop_prob, self.training, self.scale_by_keep) def extra_repr(self): return f"drop_prob={round(self.drop_prob, 3):0.3f}" class Block(nn.Module): def __init__( self, + attention, embed_dim=128, num_heads=8, ffn_ratio=4, dropout=0.1, attn_dropout=0.1, activation_dropout=0.1, activation="gelu", layer_scale_init_values=None, drop_path_rate=0.0, scale_attn_mask=False, scale_attn=True, scale_fc=True, scale_heads=True, scale_resids=True, ): super().__init__() self.embed_dim = embed_dim self.num_heads = num_heads self.head_dim = embed_dim // num_heads self.ffn_dim = embed_dim * ffn_ratio self.pre_attn_norm = nn.LayerNorm(embed_dim) - self.attn = Attention(embed_dim, num_heads, dropout=attn_dropout) + self.attn = Attention(attention, embed_dim, num_heads, dropout=attn_dropout) self.post_attn_norm = nn.LayerNorm(embed_dim) if scale_attn else nn.Identity() self.dropout = nn.Dropout(dropout) self.ls1 = ( LayerScale(embed_dim, init_values=layer_scale_init_values) if layer_scale_init_values else nn.Identity() ) self.drop_path1 = DropPath(drop_path_rate) if drop_path_rate > 0.0 else nn.Identity() self.pre_fc_norm = nn.LayerNorm(embed_dim) self.fc1 = nn.Linear(embed_dim, self.ffn_dim) if activation == "swiglu": self.fc1_g = nn.Linear(embed_dim, self.ffn_dim) self.act = nn.SiLU() else: self.fc1_g = None self.act = nn.GELU() if activation == "gelu" else nn.ReLU() self.act_dropout = nn.Dropout(activation_dropout) self.post_fc_norm = nn.LayerNorm(self.ffn_dim) if scale_fc else nn.Identity() self.fc2 = nn.Linear(self.ffn_dim, embed_dim) self.ls2 = ( LayerScale(embed_dim, init_values=layer_scale_init_values) if layer_scale_init_values else nn.Identity() ) self.drop_path2 = DropPath(drop_path_rate) if drop_path_rate > 0.0 else nn.Identity() self.c_mask = nn.Parameter(torch.ones(1), requires_grad=True) if scale_attn_mask else None self.c_attn = ( nn.Parameter(torch.ones(num_heads), requires_grad=True) if scale_heads else None ) self.w_resid = ( nn.Parameter(torch.ones(embed_dim), requires_grad=True) if scale_resids else None ) def forward(self, x, x_cls=None, padding_mask=None, attn_mask=None): """ Args: x (Tensor): input to the layer of shape `(batch, seq_len, embed_dim)` x_cls (Tensor, optional): class token input to the layer of shape `(batch, 1, embed_dim)` padding_mask (ByteTensor, optional): binary ByteTensor of shape `(batch, seq_len)` where padding elements are indicated by ``True``. Returns: encoded output of shape `(batch, seq_len, embed_dim)` """ if x_cls is not None: with torch.no_grad(): # prepend one element for x_cls: -> (batch, 1+seq_len) padding_mask = torch.cat( (torch.zeros_like(padding_mask[:, :1]), padding_mask), dim=1 ) # class attention: https://arxiv.org/pdf/2103.17239.pdf residual = x_cls u = torch.cat((x_cls, x), dim=1) # (batch, 1+seq_len, embed_dim) u = self.pre_attn_norm(u) + # default attention for convenience (could be more fancy here) x = self.attn( x_cls, u, u, key_padding_mask=padding_mask, )[0] # (1, batch, embed_dim) else: if self.c_mask is not None and attn_mask is not None: attn_mask = torch.mul(self.c_mask, attn_mask) residual = x x = self.pre_attn_norm(x) x = self.attn(x, x, x, key_padding_mask=padding_mask, attn_mask=attn_mask)[ 0 ] # (batch, seq_len, embed_dim) if self.c_attn is not None: bsz, tgt_len, _ = x.size() x = x.view(bsz, tgt_len, self.num_heads, self.head_dim) x = torch.einsum("bthd,h->btdh", x, self.c_attn) x = x.reshape(bsz, tgt_len, self.embed_dim) x = self.post_attn_norm(x) x = self.dropout(x) x = self.drop_path1(self.ls1(x)) x += residual residual = x x = self.pre_fc_norm(x) if self.fc1_g is None: x = self.act(self.fc1(x)) else: x_gate = self.fc1_g(x) x = self.fc1(x) x = self.act(x_gate) * x x = self.act_dropout(x) x = self.post_fc_norm(x) x = self.fc2(x) x = self.dropout(x) x = self.drop_path2(self.ls2(x)) if self.w_resid is not None: residual = torch.mul(self.w_resid, residual) x += residual return x class ParticleTransformer(nn.Module): """Particle Transformer (ParT) with local frame transformations.""" def __init__( self, input_dim, + attn_reps, num_classes=None, # network configurations pair_input_type="pp", pair_input_dim=None, pair_extra_dim=0, remove_self_pair=False, use_pre_activation_pair=True, embed_dims=(128, 512, 128), pair_embed_dims=(64, 64, 64), num_heads=8, num_layers=8, num_cls_layers=2, block_params=None, cls_block_params=None, fc_params=(), activation="gelu", # misc version=1, weight_init="moco", fix_init=True, trim=True, for_inference=False, for_segmentation=False, use_amp=False, **kwargs, ) -> None: super().__init__(**kwargs) self.trimmer = SequenceTrimmer(enabled=trim and not for_inference) self.for_inference = for_inference self.for_segmentation = for_segmentation self.use_amp = use_amp self.embed_dim = embed_dims[-1] if len(embed_dims) > 0 else input_dim + attn_reps = TensorReps(attn_reps) + assert attn_reps.dim * num_heads == self.embed_dim + self.attention = LLoCaAttention(attn_reps, num_heads) default_cfg = dict( embed_dim=self.embed_dim, num_heads=num_heads, ffn_ratio=4, dropout=0.1, attn_dropout=0.1, activation_dropout=0.1, activation=activation, layer_scale_init_values=None, drop_path_rate=0.0, scale_attn_mask=False, scale_fc=True, scale_attn=True, scale_heads=True, scale_resids=True, ) if version > 1: default_cfg.update( activation="swiglu", scale_fc=False, scale_attn=False, scale_heads=False, scale_resids=False, ) cfg_block = copy.deepcopy(default_cfg) if block_params is not None: cfg_block.update(block_params) cfg_cls_block = copy.deepcopy(default_cfg) cfg_cls_block.update({"dropout": 0, "attn_dropout": 0, "activation_dropout": 0}) if cls_block_params is not None: cfg_cls_block.update(cls_block_params) self.embed = ( Embed(input_dim, embed_dims, activation=activation) if len(embed_dims) > 0 else nn.Identity() ) if pair_input_dim is None: pair_input_dim = 4 if pair_input_type == "pp" else 6 self.pair_extra_dim = pair_extra_dim self.pair_embed = ( PairEmbed( pair_input_dim, pair_extra_dim, (*pair_embed_dims, cfg_block["num_heads"]), pairwise_lv_type=pair_input_type, remove_self_pair=remove_self_pair, use_pre_activation_pair=use_pre_activation_pair, for_onnx=for_inference, ) if pair_embed_dims is not None and pair_input_dim + pair_extra_dim > 0 else None ) self.blocks = nn.ModuleList( - [Block(**cfg_block) for _ in range(num_layers)] + [Block(attention=self.attention, **cfg_block) for _ in range(num_layers)] ) self.cls_blocks = ( - nn.ModuleList([Block(**cfg_cls_block) for _ in range(num_cls_layers)]) + nn.ModuleList([Block(attention=None, **cfg_cls_block) for _ in range(num_cls_layers)]) if num_cls_layers > 0 else None ) self.norm = nn.LayerNorm(self.embed_dim) if fc_params is not None: fcs = [] in_dim = self.embed_dim for param in fc_params: try: out_dim, drop_rate, act = param except ValueError: (out_dim, drop_rate), act = param, "relu" if act == "swiglu": layer = nn.Sequential( SwiGLUFFN(in_dim, out_dim * 4, out_dim, drop=drop_rate), nn.LayerNorm(out_dim), ) else: layer = nn.Sequential( nn.Linear(in_dim, out_dim), nn.GELU() if act == "gelu" else nn.ReLU(), nn.Dropout(drop_rate), ) fcs.append(layer) in_dim = out_dim fcs.append(nn.Linear(in_dim, num_classes)) self.fc = nn.Sequential(*fcs) else: self.fc = None # cls tokens if not self.for_segmentation and num_cls_layers > 0: self.cls_token = nn.Parameter(torch.zeros(1, 1, self.embed_dim), requires_grad=True) nn.init.trunc_normal_(self.cls_token, std=0.02) else: self.cls_token = None # weight initialization if weight_init is not None: self.init_weights(weight_init) if fix_init: self.fix_init_weight() def fix_init_weight(self): def rescale(param, _layer_id): param.div_(math.sqrt(2.0 * _layer_id)) for layer_id, layer in enumerate(self.blocks): rescale(layer.attn.out_proj.weight.data, layer_id + 1) rescale(layer.fc2.weight.data, layer_id + 1) def init_weights(self, mode: str = "") -> None: assert mode in ("timm", "moco") if mode == "timm": named_apply(init_weights_vit_timm, self) elif mode == "moco": named_apply(init_weights_vit_moco, self) @torch.jit.ignore def no_weight_decay(self): return { "cls_token", } def _forward_encoder(self, x, v=None, mask=None, uu=None, uu_idx=None): with torch.no_grad(): if not self.for_inference: if uu_idx is not None: uu = build_sparse_tensor(uu, uu_idx, x.size(-1)) x, v, mask, uu = self.trimmer(x, v, mask, uu) padding_mask = ~mask.squeeze(1) # (batch_size, seq_len) with torch.autocast("cuda", enabled=self.use_amp): # input embedding x = self.embed(x).masked_fill( ~mask.transpose(1, 2), 0 ) # (batch_size, seq_len, num_fts) attn_mask = None if (v is not None or uu is not None) and self.pair_embed is not None: attn_mask = self.pair_embed( v, uu=uu, mask=mask ) # (batch_size, num_heads, seq_len, seq_len) # transform for block in self.blocks: x = block(x, x_cls=None, padding_mask=padding_mask, attn_mask=attn_mask) # x: (batch, seq_len, embed_dim) # padding_mask: (batch, seq_len) return x, padding_mask def _forward_aggregator(self, x, padding_mask): with torch.autocast("cuda", enabled=self.use_amp): if self.cls_blocks is not None: # for classification: extract using class token cls_tokens = self.cls_token.expand(x.size(0), 1, -1) # (batch, 1, embed_dim) for block in self.cls_blocks: cls_tokens = block( x, x_cls=cls_tokens, padding_mask=padding_mask ) # (batch, 1, embed_dim) cls_tokens = cls_tokens.squeeze(1) # (batch, embed_dim) else: # for classification: simple average pooling mask = ~padding_mask.unsqueeze(1) # (batch, 1, seq_len) x = x.transpose(1, 2).contiguous() # (batch, embed_dim, seq_len) counts = mask.float().sum(-1) # (batch, 1) counts = torch.max(counts, torch.ones_like(counts)) # >=1 cls_tokens = (x * mask).sum(-1) / counts # (batch, embed_dim) x_cls = self.norm(cls_tokens) # (batch, embed_dim) return x_cls - def forward(self, x, v=None, mask=None, uu=None, uu_idx=None): + def forward(self, x, frames, v=None, mask=None, uu=None, uu_idx=None): # x: (batch_size, num_fts, seq_len) # v: (batch_size, 4, seq_len) [px,py,pz,energy] # mask: (batch_size, 1, seq_len) -- real particle = 1, padded = 0 # for pytorch: uu (batch_size, C', num_pairs), uu_idx (batch_size, 2, num_pairs) # for onnx: uu (batch_size, C', seq_len, seq_len), uu_idx=None + self.attention.prepare_frames(frames) x, padding_mask = self._forward_encoder(x, v=v, mask=mask, uu=uu, uu_idx=uu_idx) if self.cls_blocks is None and self.fc is None: # x: (batch, seq_len, embed_dim) # padding_mask: (batch, seq_len) return x, padding_mask with torch.autocast("cuda", enabled=self.use_amp): # === for segmentation === if self.for_segmentation: x = self.norm(x) if self.fc is not None: x = self.fc(x) # x: (P, N, C) -> output: (N, C, P) output = x.transpose(1, 2).contiguous() if self.for_inference: output = torch.softmax(output, dim=1) # print('output:\n', output) return output x_cls = self._forward_aggregator(x, padding_mask) if self.fc is None: return x_cls # fc output = self.fc(x_cls) if self.for_inference: output = torch.softmax(output, dim=1) # print('output:\n', output) return output ### weight initialization methods ### def init_weights_vit_timm(module: nn.Module, name: str = "") -> None: """ViT weight initialization, original timm impl (for reproducibility)""" if isinstance(module, nn.Linear): nn.init.trunc_normal_(module.weight, std=0.02) if module.bias is not None: nn.init.zeros_(module.bias) elif hasattr(module, "init_weights"): module.init_weights() def init_weights_vit_moco(module: nn.Module, name: str = "") -> None: """ViT weight initialization, matching moco-v3 impl minus fixed PatchEmbed""" if isinstance(module, nn.Linear): if "in_proj" in name: # treat the weights of Q, K, V separately val = math.sqrt(6.0 / float(module.weight.shape[0] // 3 + module.weight.shape[1])) nn.init.uniform_(module.weight, -val, val) else: nn.init.xavier_uniform_(module.weight) if module.bias is not None: nn.init.zeros_(module.bias) elif hasattr(module, "init_weights"): module.init_weights() def named_apply( fn: Callable, module: nn.Module, name="", depth_first: bool = True, include_root: bool = False, ) -> nn.Module: if not depth_first and include_root: fn(module=module, name=name) for child_name, child_module in module.named_children(): child_name = ".".join((name, child_name)) if name else child_name named_apply( fn=fn, module=child_module, name=child_name, depth_first=depth_first, include_root=True, ) if depth_first and include_root: fn(module=module, name=name) return module