wehub-resource-sync
58 lines
1.6 KiB
Python
58 lines
1.6 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
from dgl.nn.pytorch import GraphConv
|
|
from train import device
|
|
|
|
|
|
class VGAEModel(nn.Module):
|
|
def __init__(self, in_dim, hidden1_dim, hidden2_dim):
|
|
super(VGAEModel, self).__init__()
|
|
self.in_dim = in_dim
|
|
self.hidden1_dim = hidden1_dim
|
|
self.hidden2_dim = hidden2_dim
|
|
|
|
layers = [
|
|
GraphConv(
|
|
self.in_dim,
|
|
self.hidden1_dim,
|
|
activation=F.relu,
|
|
allow_zero_in_degree=True,
|
|
),
|
|
GraphConv(
|
|
self.hidden1_dim,
|
|
self.hidden2_dim,
|
|
activation=lambda x: x,
|
|
allow_zero_in_degree=True,
|
|
),
|
|
GraphConv(
|
|
self.hidden1_dim,
|
|
self.hidden2_dim,
|
|
activation=lambda x: x,
|
|
allow_zero_in_degree=True,
|
|
),
|
|
]
|
|
self.layers = nn.ModuleList(layers)
|
|
|
|
def encoder(self, g, features):
|
|
h = self.layers[0](g, features)
|
|
self.mean = self.layers[1](g, h)
|
|
self.log_std = self.layers[2](g, h)
|
|
gaussian_noise = torch.randn(features.size(0), self.hidden2_dim).to(
|
|
device
|
|
)
|
|
sampled_z = self.mean + gaussian_noise * torch.exp(self.log_std).to(
|
|
device
|
|
)
|
|
return sampled_z
|
|
|
|
def decoder(self, z):
|
|
adj_rec = torch.sigmoid(torch.matmul(z, z.t()))
|
|
return adj_rec
|
|
|
|
def forward(self, g, features):
|
|
z = self.encoder(g, features)
|
|
adj_rec = self.decoder(z)
|
|
return adj_rec
|