Spaces:
Runtime error
Runtime error
| ################################################################### | |
| # 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 | |