File size: 2,569 Bytes
e87a462 |
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 |
from torch_geometric.nn.conv import MessagePassing
from torch_geometric.nn.conv.cheb_conv import ChebConv
from torch_geometric.nn.inits import zeros, normal
# We change the default initialization from zeros to a normal distribution
class ChebConv(ChebConv):
def reset_parameters(self):
for lin in self.lins:
normal(lin, mean = 0, std = 0.1)
#lin.reset_parameters()
normal(self.bias, mean = 0, std = 0.1)
#zeros(self.bias)
# Pooling from COMA: https://github.com/pixelite1201/pytorch_coma/blob/master/layers.py
class Pool(MessagePassing):
def __init__(self):
# source_to_target is the default value for flow, but is specified here for explicitness
super(Pool, self).__init__(flow='source_to_target')
def forward(self, x, pool_mat, dtype=None):
pool_mat = pool_mat.transpose(0, 1)
out = self.propagate(edge_index=pool_mat._indices(), x=x, norm=pool_mat._values(), size=pool_mat.size())
return out
def message(self, x_j, norm):
return norm.view(1, -1, 1) * x_j
import torch.nn as nn
import torch.nn.functional as F
class residualBlock(nn.Module):
def __init__(self, in_channels, out_channels, stride=1):
"""
Args:
in_channels (int): Number of input channels.
out_channels (int): Number of output channels.
stride (int): Controls the stride.
"""
super(residualBlock, self).__init__()
self.skip = nn.Sequential()
if stride != 1 or in_channels != out_channels:
self.skip = nn.Sequential(
nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride=stride, bias=False),
nn.BatchNorm2d(out_channels, track_running_stats=False))
else:
self.skip = None
self.block = nn.Sequential(nn.BatchNorm2d(in_channels, track_running_stats=False),
nn.ReLU(inplace=True),
nn.Conv2d(in_channels, out_channels, 3, padding=1),
nn.BatchNorm2d(out_channels, track_running_stats=False),
nn.ReLU(inplace=True),
nn.Conv2d(out_channels, out_channels, 3, padding=1)
)
def forward(self, x):
identity = x
out = self.block(x)
if self.skip is not None:
identity = self.skip(x)
out += identity
out = F.relu(out)
return out |