| | import math |
| | from typing import Optional, Callable, List |
| |
|
| | import numpy as np |
| | import torch |
| | import torch.nn as nn |
| | from scipy.stats import truncnorm |
| |
|
| |
|
| | def _prod(nums): |
| | out = 1 |
| | for n in nums: |
| | out = out * n |
| | return out |
| |
|
| | def _calculate_fan(linear_weight_shape, fan="fan_in"): |
| | fan_out, fan_in = linear_weight_shape |
| |
|
| | if fan == "fan_in": |
| | f = fan_in |
| | elif fan == "fan_out": |
| | f = fan_out |
| | elif fan == "fan_avg": |
| | f = (fan_in + fan_out) / 2 |
| | else: |
| | raise ValueError("Invalid fan option") |
| |
|
| | return f |
| |
|
| | def trunc_normal_init_(weights, scale=1.0, fan="fan_in"): |
| | shape = weights.shape |
| | f = _calculate_fan(shape, fan) |
| | scale = scale / max(1, f) |
| | a = -2 |
| | b = 2 |
| | std = math.sqrt(scale) / truncnorm.std(a=a, b=b, loc=0, scale=1) |
| | size = _prod(shape) |
| | samples = truncnorm.rvs(a=a, b=b, loc=0, scale=std, size=size) |
| | samples = np.reshape(samples, shape) |
| | with torch.no_grad(): |
| | weights.copy_(torch.tensor(samples, device=weights.device)) |
| |
|
| | def lecun_normal_init_(weights): |
| | trunc_normal_init_(weights, scale=1.0) |
| |
|
| | def he_normal_init_(weights): |
| | trunc_normal_init_(weights, scale=2.0) |
| |
|
| | def glorot_uniform_init_(weights): |
| | nn.init.xavier_uniform_(weights, gain=1) |
| |
|
| | def final_init_(weights): |
| | with torch.no_grad(): |
| | weights.fill_(0.0) |
| |
|
| | def gating_init_(weights): |
| | with torch.no_grad(): |
| | weights.fill_(0.0) |
| |
|
| | def normal_init_(weights): |
| | nn.init.kaiming_normal_(weights, nonlinearity="linear") |
| |
|
| |
|
| | class Linear(nn.Linear): |
| | """ |
| | A Linear layer with in-house nonstandard initializations. |
| | """ |
| |
|
| | def __init__( |
| | self, |
| | in_dim: int, |
| | out_dim: int, |
| | bias: bool = True, |
| | init: str = "default", |
| | init_fn: Optional[Callable[[torch.Tensor, torch.Tensor], None]] = None, |
| | ): |
| | """ |
| | Args: |
| | in_dim: |
| | The final dimension of inputs to the layer |
| | out_dim: |
| | The final dimension of layer outputs |
| | bias: |
| | Whether to learn an additive bias. True by default |
| | init: |
| | The initializer to use. Choose from: |
| | |
| | "default": LeCun fan-in truncated normal initialization |
| | "relu": He initialization w/ truncated normal distribution |
| | "glorot": Fan-average Glorot uniform initialization |
| | "gating": Weights=0, Bias=1 |
| | "normal": Normal initialization with std=1/sqrt(fan_in) |
| | "final": Weights=0, Bias=0 |
| | |
| | Overridden by init_fn if the latter is not None. |
| | init_fn: |
| | A custom initializer taking weight and bias as inputs. |
| | Overrides init if not None. |
| | """ |
| | super(Linear, self).__init__(in_dim, out_dim, bias=bias) |
| |
|
| | if bias: |
| | with torch.no_grad(): |
| | self.bias.fill_(0) |
| | if init_fn is not None: |
| | init_fn(self.weight, self.bias) |
| | else: |
| | if init == "default": |
| | lecun_normal_init_(self.weight) |
| | elif init == "relu": |
| | he_normal_init_(self.weight) |
| | elif init == "glorot": |
| | glorot_uniform_init_(self.weight) |
| | elif init == "gating": |
| | gating_init_(self.weight) |
| | if bias: |
| | with torch.no_grad(): |
| | self.bias.fill_(1.0) |
| | elif init == "normal": |
| | normal_init_(self.weight) |
| | elif init == "final": |
| | final_init_(self.weight) |
| | else: |
| | raise ValueError("Invalid init string.") |
| |
|
| |
|
| | |
| | class NodeTransition(nn.Module): |
| | def __init__(self, dim): |
| | super(NodeTransition, self).__init__() |
| | self.dim = dim |
| | self.linear_1 = Linear(dim, dim, init="relu") |
| | self.linear_2 = Linear(dim, dim, init="relu") |
| | self.linear_3 = Linear(dim, dim, init="final") |
| | self.relu = nn.ReLU() |
| | self.ln = nn.LayerNorm(dim) |
| |
|
| | def forward(self, s): |
| | s_initial = s |
| | s = self.relu(self.linear_1(s)) |
| | s = self.relu(self.linear_2(s)) |
| | s = self.linear_3(s) |
| | s = s + s_initial |
| | s = self.ln(s) |
| | return s |
| |
|
| |
|
| | class EdgeTransition(nn.Module): |
| | def __init__(self, |
| | node_embed_size, |
| | edge_embed_in, |
| | edge_embed_out, |
| | num_layers=2, |
| | node_dilation=2 |
| | ): |
| | super(EdgeTransition, self).__init__() |
| |
|
| | bias_embed_size = node_embed_size // node_dilation |
| | self.initial_embed = Linear( |
| | node_embed_size, bias_embed_size, init="relu") |
| | hidden_size = bias_embed_size * 2 + edge_embed_in |
| | trunk_layers = [] |
| | for _ in range(num_layers): |
| | trunk_layers.append(Linear(hidden_size, hidden_size, init="relu")) |
| | trunk_layers.append(nn.ReLU()) |
| | self.trunk = nn.Sequential(*trunk_layers) |
| | self.final_layer = Linear(hidden_size, edge_embed_out, init="final") |
| | self.layer_norm = nn.LayerNorm(edge_embed_out) |
| |
|
| | def forward(self, node_embed, edge_embed): |
| | node_embed = self.initial_embed(node_embed) |
| | batch_size, num_res, _ = node_embed.shape |
| | edge_bias = torch.cat([ |
| | torch.tile(node_embed[:, :, None, :], (1, 1, num_res, 1)), |
| | torch.tile(node_embed[:, None, :, :], (1, num_res, 1, 1)), |
| | ], axis=-1) |
| | edge_embed = torch.cat( |
| | [edge_embed, edge_bias], axis=-1).reshape( |
| | batch_size * num_res**2, -1) |
| | edge_embed = self.final_layer(self.trunk(edge_embed) + edge_embed) |
| | edge_embed = self.layer_norm(edge_embed) |
| | edge_embed = edge_embed.reshape( |
| | batch_size, num_res, num_res, -1 |
| | ) |
| | return edge_embed |
| | |
| | |
| | class TorsionAngleHead(nn.Module): |
| | def __init__(self, in_dim, n_torsion_angles, eps=1e-8): |
| | super(TorsionAngleHead, self).__init__() |
| | |
| | self.linear_1 = Linear(in_dim, in_dim, init="relu") |
| | self.linear_2 = Linear(in_dim, in_dim, init="relu") |
| | self.linear_3 = Linear(in_dim, in_dim, init="final") |
| | self.linear_final = Linear(in_dim, n_torsion_angles * 2, init="final") |
| | self.relu = nn.ReLU() |
| | self.eps = eps |
| |
|
| | def forward(self, s): |
| | s_initial = s |
| | s = self.relu(self.linear_1(s)) |
| | s = self.linear_2(s) |
| | s = s + s_initial |
| | |
| | unnormalized_s = self.linear_final(s) |
| | norm_denom = torch.sqrt( |
| | torch.clamp( |
| | torch.sum(unnormalized_s ** 2, dim=-1, keepdim=True), |
| | min=self.eps, |
| | ) |
| | ) |
| | normalized_s = unnormalized_s / norm_denom |
| | return normalized_s |
| | |
| | |
| | class BackboneUpdate(nn.Module): |
| | """ |
| | Implements part of Algorithm 23. |
| | """ |
| |
|
| | def __init__(self, c_s): |
| | """ |
| | Args: |
| | c_s: |
| | Single representation channel dimension |
| | """ |
| | super(BackboneUpdate, self).__init__() |
| |
|
| | self.c_s = c_s |
| | self.linear = Linear(self.c_s, 6, init="final") |
| |
|
| | def forward(self, s: torch.Tensor): |
| | """ |
| | Args: |
| | [*, N_res, C_s] single representation |
| | Returns: |
| | [*, N_res, 6] update vector |
| | """ |
| | |
| | update = self.linear(s) |
| | return update |
