wehub-resource-sync
18 lines
498 B
Python
Executable File
18 lines
498 B
Python
Executable File
import torch
|
|
from torch import nn as nn
|
|
from torch.autograd import Variable
|
|
from torch.nn import functional as F
|
|
|
|
from model.tensorized_layers.graphsage import BatchedGraphSAGE
|
|
|
|
|
|
class DiffPoolAssignment(nn.Module):
|
|
def __init__(self, nfeat, nnext):
|
|
super().__init__()
|
|
self.assign_mat = BatchedGraphSAGE(nfeat, nnext, use_bn=True)
|
|
|
|
def forward(self, x, adj, log=False):
|
|
s_l_init = self.assign_mat(x, adj)
|
|
s_l = F.softmax(s_l_init, dim=-1)
|
|
return s_l
|