| 86 | |
| 87 | |
| 88 | class GrokFormer(nn.Module): |
| 89 | |
| 90 | def __init__(self, nclass, nfeat, nlayer=1, hidden_dim=128, dim=32, nheads=1, k=10, |
| 91 | tran_dropout=0.0, feat_dropout=0.0, prop_dropout=0.0, norm='none'): |
| 92 | super(GrokFormer, self).__init__() |
| 93 | |
| 94 | self.norm = norm |
| 95 | self.nfeat = nfeat |
| 96 | self.nlayer = nlayer |
| 97 | self.nheads = nheads |
| 98 | self.hidden_dim = hidden_dim |
| 99 | self.dim = dim |
| 100 | |
| 101 | self.feat_encoder = nn.Sequential( |
| 102 | nn.Linear(nfeat, hidden_dim), |
| 103 | nn.ReLU(), |
| 104 | nn.Linear(hidden_dim, nclass), |
| 105 | ) |
| 106 | |
| 107 | |
| 108 | self.eig_encoder = SineEncoding(k, dim) |
| 109 | |
| 110 | self.mha_dropout = nn.Dropout(tran_dropout) |
| 111 | self.ffn_dropout = nn.Dropout(tran_dropout) |
| 112 | self.prop_dropout = nn.Dropout(prop_dropout) |
| 113 | |
| 114 | self.k = k |
| 115 | self.alpha = nn.Linear(self.k, 1, bias=False) |
| 116 | |
| 117 | self.feat_dp1 = nn.Dropout(feat_dropout) |
| 118 | self.feat_dp2 = nn.Dropout(feat_dropout) |
| 119 | |
| 120 | if norm == 'none': |
| 121 | self.mha_norm = nn.LayerNorm(nclass) |
| 122 | self.ffn_norm = nn.LayerNorm(nclass) |
| 123 | self.mha = MultiHeadAttention(nclass, nheads, tran_dropout) |
| 124 | self.ffn = FeedForwardNetwork(nclass, nclass, nclass) |
| 125 | else: |
| 126 | self.mha_norm = nn.LayerNorm(hidden_dim) |
| 127 | self.ffn_norm = nn.LayerNorm(hidden_dim) |
| 128 | self.mha = MultiHeadAttention(hidden_dim, nheads, tran_dropout) |
| 129 | self.ffn = FeedForwardNetwork(hidden_dim, hidden_dim, hidden_dim) |
| 130 | self.classify = nn.Linear(hidden_dim, nclass) |
| 131 | |
| 132 | def transformer_encoder(self, h, h_fur): |
| 133 | mha_h = self.mha_norm(h) |
| 134 | mha_h = self.mha(mha_h, mha_h, mha_h) |
| 135 | mha_h_ = h + self.mha_dropout(mha_h) + h_fur |
| 136 | |
| 137 | ffn_h = self.ffn_norm(mha_h_) |
| 138 | ffn_h = self.ffn(ffn_h) |
| 139 | encoder_h = mha_h_ + self.ffn_dropout(ffn_h) |
| 140 | return encoder_h |
| 141 | |
| 142 | def forward(self, e, u, x): |
| 143 | N = e.size(0) |
| 144 | ut = u.permute(1, 0) |
| 145 | |