Spaces:
Runtime error
Runtime error
File size: 2,284 Bytes
01bb3bb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 |
###################################################################
# File Name: GCN.py
# Author: S.X.Zhang
###################################################################
from __future__ import print_function
from __future__ import division
from __future__ import absolute_import
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn import init
class MeanAggregator(nn.Module):
def __init__(self):
super(MeanAggregator, self).__init__()
def forward(self, features, A):
x = torch.bmm(A, features)
return x
class GraphConv(nn.Module):
def __init__(self, in_dim, out_dim, agg):
super(GraphConv, self).__init__()
self.in_dim = in_dim
self.out_dim = out_dim
self.weight = nn.Parameter(torch.FloatTensor(in_dim * 2, out_dim))
self.bias = nn.Parameter(torch.FloatTensor(out_dim))
init.xavier_uniform_(self.weight)
init.constant_(self.bias, 0)
self.agg = agg()
def forward(self, features, A):
b, n, d = features.shape
assert (d == self.in_dim)
agg_feats = self.agg(features, A)
cat_feats = torch.cat([features, agg_feats], dim=2)
out = torch.einsum('bnd,df->bnf', (cat_feats, self.weight))
out = F.relu(out + self.bias)
return out
class GCN(nn.Module):
def __init__(self, in_dim, out_dim):
super(GCN, self).__init__()
self.bn0 = nn.BatchNorm1d(in_dim, affine=False)
self.conv1 = GraphConv(in_dim, 256, MeanAggregator)
self.conv2 = GraphConv(256, 1024, MeanAggregator)
self.conv3 = GraphConv(1024, 512, MeanAggregator)
self.conv4 = GraphConv(512, out_dim, MeanAggregator)
self.prediction = nn.Sequential(
nn.Conv1d(out_dim, 128, 1),
nn.ReLU(inplace=True),
nn.Conv1d(128, 64, 1),
nn.ReLU(inplace=True),
nn.Conv1d(64, 2, 1))
def forward(self, x, A):
x = self.bn0(x)
x = x.permute(0, 2, 1)
b, n, c = x.shape
A = A.expand(b, n, n)
x = self.conv1(x, A)
x = self.conv2(x, A)
x = self.conv3(x, A)
x = self.conv4(x, A)
x = x.permute(0, 2, 1)
pred = self.prediction(x)
return pred
|