(self, d_msa=256, d_pair=128, d_state=16,
SE3_param={'l0_in_features':32, 'l0_out_features':16, 'num_edge_features':32}, p_drop=0.1)
| 200 | |
| 201 | class Str2Str(nn.Module): |
| 202 | def __init__(self, d_msa=256, d_pair=128, d_state=16, |
| 203 | SE3_param={'l0_in_features':32, 'l0_out_features':16, 'num_edge_features':32}, p_drop=0.1): |
| 204 | super(Str2Str, self).__init__() |
| 205 | |
| 206 | # initial node & pair feature process |
| 207 | self.norm_msa = nn.LayerNorm(d_msa) |
| 208 | self.norm_pair = nn.LayerNorm(d_pair) |
| 209 | self.norm_state = nn.LayerNorm(d_state) |
| 210 | |
| 211 | self.embed_x = nn.Linear(d_msa+d_state, SE3_param['l0_in_features']) |
| 212 | self.embed_e1 = nn.Linear(d_pair, SE3_param['num_edge_features']) |
| 213 | self.embed_e2 = nn.Linear(SE3_param['num_edge_features']+36+1, SE3_param['num_edge_features']) |
| 214 | |
| 215 | self.norm_node = nn.LayerNorm(SE3_param['l0_in_features']) |
| 216 | self.norm_edge1 = nn.LayerNorm(SE3_param['num_edge_features']) |
| 217 | self.norm_edge2 = nn.LayerNorm(SE3_param['num_edge_features']) |
| 218 | |
| 219 | self.se3 = SE3TransformerWrapper(**SE3_param) |
| 220 | self.sc_predictor = SCPred(d_msa=d_msa, d_state=SE3_param['l0_out_features'], |
| 221 | p_drop=p_drop) |
| 222 | |
| 223 | self.reset_parameter() |
| 224 | |
| 225 | def reset_parameter(self): |
| 226 | # initialize weights to normal distribution |
nothing calls this directly
no test coverage detected