dmlc--dgl
Hongzhi (Steve), Chen
23d09057e3
* [Misc] Black auto fix. * sort Co-authored-by: Steve <ubuntu@ip-172-31-34-29.ap-northeast-1.compute.internal>
24 行
616 B
Python
24 行
616 B
Python
import torch
|
|
import torch.nn as nn
|
|
|
|
|
|
class EntropyLoss(nn.Module):
|
|
# Return Scalar
|
|
def forward(self, adj, anext, s_l):
|
|
entropy = (
|
|
(torch.distributions.Categorical(probs=s_l).entropy())
|
|
.sum(-1)
|
|
.mean(-1)
|
|
)
|
|
assert not torch.isnan(entropy)
|
|
return entropy
|
|
|
|
|
|
class LinkPredLoss(nn.Module):
|
|
def forward(self, adj, anext, s_l):
|
|
link_pred_loss = (adj - s_l.matmul(s_l.transpose(-1, -2))).norm(
|
|
dim=(1, 2)
|
|
)
|
|
link_pred_loss = link_pred_loss / (adj.size(1) * adj.size(2))
|
|
return link_pred_loss.mean()
|