def __init__(self, in_features, out_features, aggregators, scalers, avg_d,
self_loop, pretrans_layers, posttrans_layers, device):
"""
:param in_features: size of the input per node of the tower
:param out_features: size of the output per node of the tower
:param aggregators: set of aggregation functions each taking as input X (B x N x N x Din), adj (B x N x N), self_loop and device
:param scalers: set of scaling functions each taking as input X (B x N x Din), adj (B x N x N) and avg_d
"""
super(PNATower, self).__init__()
self.device = device
self.in_features = in_features
self.out_features = out_features
self.aggregators = aggregators
self.scalers = scalers
self.self_loop = self_loop
self.pretrans = MLP(in_size=2 * self.in_features,
hidden_size=self.in_features,
out_size=self.in_features,
layers=pretrans_layers,
mid_activation='relu',
last_activation='none')
self.posttrans = MLP(in_size=(len(aggregators) * len(scalers) + 1) *
self.in_features,
hidden_size=self.out_features,
out_size=self.out_features,
layers=posttrans_layers,
mid_activation='relu',
last_activation='none')
self.avg_d = avg_d
def __init__(self, in_dim, out_dim, dropout, graph_norm, batch_norm,
aggregators, scalers, avg_d, pretrans_layers,
posttrans_layers, edge_features, edge_dim):
super().__init__()
self.dropout = dropout
self.graph_norm = graph_norm
self.batch_norm = batch_norm
self.edge_features = edge_features
self.fc = nn.Linear(in_dim, out_dim, bias=False)
self.attn_fc = nn.Linear(2 * out_dim, 1, bias=False)
self.batchnorm_h = nn.BatchNorm1d(out_dim)
self.aggregators = aggregators
self.scalers = scalers
self.pretrans = MLP(in_size=2 * in_dim +
(edge_dim if edge_features else 0),
hidden_size=in_dim,
out_size=in_dim,
layers=pretrans_layers,
mid_activation='relu',
last_activation='none')
self.posttrans = MLP(in_size=(len(aggregators) * len(scalers) + 1) *
in_dim,
hidden_size=out_dim,
out_size=out_dim,
layers=posttrans_layers,
mid_activation='relu',
last_activation='none')
self.avg_d = avg_d
def __init__(
self,
input_dim=2,
h_dim=64,
z_dim=2,
nonlinearity='tanh',
num_hidden_layers=1,
init='gaussian', #None,
):
super().__init__()
self.input_dim = input_dim
self.h_dim = h_dim
self.z_dim = z_dim
self.nonlinearity = nonlinearity
self.num_hidden_layers = num_hidden_layers
self.init = init
self.main = MLP(input_dim=z_dim,
hidden_dim=h_dim,
output_dim=h_dim,
nonlinearity=nonlinearity,
num_hidden_layers=num_hidden_layers - 1,
use_nonlinearity_output=True)
self.reparam = NormalDistributionLinear(h_dim, input_dim)
if self.init == 'gaussian':
self.reset_parameters()
else:
pass
def __init__(self, config):
super(TSPModel, self).__init__()
# Define net parameters
self.num_nodes = config['num_nodes']
self.node_dim = config['node_dim']
self.voc_nodes_in = config['voc_nodes_in']
self.voc_nodes_out = config['num_nodes'] # config['voc_nodes_out']
self.voc_edges_in = config['voc_edges_in']
self.voc_edges_out = config['voc_edges_out']
self.hidden_dim = config['hidden_dim']
self.num_layers = config['num_layers']
self.mlp_layers = config['mlp_layers']
# Node and edge embedding layers/lookups
self.nodes_embedding = nn.Linear(self.node_dim,
self.hidden_dim,
bias=False)
self.edges_embedding = nn.Linear(1 + self.voc_edges_in,
self.hidden_dim,
bias=False)
gcn_layers = []
for layer in range(self.num_layers):
gcn_layers.append(TSPConv(self.hidden_dim))
self.gcn_layers = nn.ModuleList(gcn_layers)
# Define MLP classifiers
self.mlp_edges = MLP(self.hidden_dim, self.voc_edges_out,
self.mlp_layers)
def __init__(
self,
input_dim=2,
h_dim=1000,
std=0.1,
num_hidden_layers=1,
nonlinearity='tanh',
noise_type='gaussian',
#init=True,
):
super().__init__()
self.input_dim = input_dim
self.h_dim = h_dim
self.std = std
self.num_hidden_layers = num_hidden_layers
self.nonlinearity = nonlinearity
self.noise_type = noise_type
#self.init = init
self.neglogprob = MLP(input_dim + 1,
h_dim,
1,
use_nonlinearity_output=False,
num_hidden_layers=num_hidden_layers,
nonlinearity=nonlinearity)
def __init__(
self,
input_dim=2,
h_dim=8,
noise_dim=2,
nonlinearity='softplus',
num_hidden_layers=1,
clip_logvar=None,
):
super().__init__()
self.input_dim = input_dim
self.h_dim = h_dim
self.noise_dim = noise_dim
self.nonlinearity = nonlinearity
self.num_hidden_layers = num_hidden_layers
self.clip_logvar = clip_logvar
self.main = MLP(input_dim=input_dim,
hidden_dim=h_dim,
output_dim=h_dim,
nonlinearity=nonlinearity,
num_hidden_layers=num_hidden_layers - 1,
use_nonlinearity_output=True)
self.reparam = NormalDistributionLinear(h_dim,
noise_dim,
nonlinearity=clip_logvar)
def __init__(
self,
input_height=28,
input_channels=1,
noise_dim=100,
z_dim=32,
nonlinearity='softplus',
enc_noise=False,
):
super().__init__()
self.input_height = input_height
self.input_channels = input_channels
self.noise_dim = noise_dim
self.z_dim = z_dim
self.nonlinearity = nonlinearity
self.enc_noise = enc_noise
h_dim = 256
nos_dim = noise_dim if not enc_noise else h_dim
s_h = input_height
s_h2 = conv_out_size(s_h, 5, 2, 2)
s_h4 = conv_out_size(s_h2, 5, 2, 2)
s_h8 = conv_out_size(s_h4, 5, 2, 2)
#print(s_h, s_h2, s_h4, s_h8)
self.afun = get_nonlinear_func(nonlinearity)
self.conv1 = nn.Conv2d(self.input_channels, 16, 5, 2, 2, bias=True)
self.conv2 = nn.Conv2d(16, 32, 5, 2, 2, bias=True)
self.conv3 = nn.Conv2d(32, 32, 5, 2, 2, bias=True)
self.fc4 = nn.Linear(s_h8 * s_h8 * 32 + nos_dim, 800, bias=True)
self.fc5 = nn.Linear(800, z_dim, bias=True)
self.nos_encode = Identity() if not enc_noise \
else MLP(input_dim=noise_dim, hidden_dim=h_dim, output_dim=h_dim, nonlinearity=nonlinearity, num_hidden_layers=2, use_nonlinearity_output=True)
def __init__(
self,
input_dim=2,
noise_dim=2,
h_dim=64,
z_dim=2,
nonlinearity='softplus',
num_hidden_layers=1,
std=1.,
init='none', #'gaussian',
enc_noise=False,
):
super().__init__(
input_dim=input_dim,
noise_dim=noise_dim,
h_dim=h_dim,
z_dim=z_dim,
nonlinearity=nonlinearity,
num_hidden_layers=num_hidden_layers,
std=std,
init=init,
enc_noise=enc_noise,
)
nos_dim = noise_dim if not enc_noise else h_dim
self.inp_encode = MLP(input_dim=input_dim,
hidden_dim=h_dim,
output_dim=h_dim,
nonlinearity=nonlinearity,
num_hidden_layers=num_hidden_layers,
use_nonlinearity_output=True)
self.nos_encode = Identity() if not enc_noise \
else MLP(input_dim=noise_dim, hidden_dim=h_dim, output_dim=h_dim, nonlinearity=nonlinearity, num_hidden_layers=0, use_nonlinearity_output=True)
self.fc = MLP(input_dim=h_dim + nos_dim,
hidden_dim=h_dim,
output_dim=z_dim,
nonlinearity=nonlinearity,
num_hidden_layers=1,
use_nonlinearity_output=False)
if self.init == 'gaussian':
self.reset_parameters()
else:
pass
def __init__(
self,
input_dim=2, #10,
h_dim=128,
context_dim=2,
std=0.01,
num_hidden_layers=1,
nonlinearity='softplus',
noise_type='gaussian',
enc_input=False, # True,
enc_ctx=False, #True,
#init=True,
std_method='default',
):
super().__init__()
self.input_dim = input_dim
self.h_dim = h_dim
self.context_dim = context_dim
self.std = std
self.num_hidden_layers = num_hidden_layers
self.nonlinearity = nonlinearity
self.noise_type = noise_type
self.enc_input = enc_input
if self.enc_input:
inp_dim = h_dim
else:
inp_dim = input_dim
self.enc_ctx = enc_ctx
if self.enc_ctx:
ctx_dim = h_dim
else:
ctx_dim = context_dim
self.ctx_dim = ctx_dim
self.ctx_encode = Identity() if not self.enc_ctx \
else MLP(context_dim, h_dim, h_dim, nonlinearity=nonlinearity, num_hidden_layers=num_hidden_layers-1 if self.enc_ctx == 'deep' else 1, use_nonlinearity_output=True)
self.inp_encode = Identity() if not self.enc_input \
else MLP(input_dim, h_dim, h_dim, nonlinearity=nonlinearity, num_hidden_layers=num_hidden_layers-1, use_nonlinearity_output=True)
self.neglogprob = MLP(inp_dim + ctx_dim + 1,
h_dim,
1,
nonlinearity=nonlinearity,
num_hidden_layers=num_hidden_layers,
use_nonlinearity_output=False)
def __init__(
self,
input_dim=2, #10,
h_dim=128,
context_dim=2,
std=0.01,
num_hidden_layers=1,
nonlinearity='tanh',
noise_type='gaussian',
enc_input=True,
enc_ctx=True,
#init=True,
):
super().__init__()
self.input_dim = input_dim
self.h_dim = h_dim
self.context_dim = context_dim
self.std = std
self.num_hidden_layers = num_hidden_layers
self.nonlinearity = nonlinearity
self.noise_type = noise_type
self.enc_input = enc_input
if self.enc_input:
inp_dim = h_dim
else:
inp_dim = input_dim
self.enc_ctx = enc_ctx
if self.enc_ctx:
ctx_dim = h_dim
else:
ctx_dim = context_dim
#self.init = init
self.ctx_encode = Identity() if not self.enc_ctx \
else MLP(context_dim, h_dim, h_dim, nonlinearity=nonlinearity, num_hidden_layers=num_hidden_layers-1, use_nonlinearity_output=True)
self.inp_encode = Identity() if not self.enc_input \
else MLP(input_dim, h_dim, h_dim, nonlinearity=nonlinearity, num_hidden_layers=num_hidden_layers-1, use_nonlinearity_output=True)
self.dae = MLP(inp_dim + ctx_dim,
h_dim,
input_dim,
nonlinearity=nonlinearity,
num_hidden_layers=num_hidden_layers,
use_nonlinearity_output=False)
def __init__(
self,
input_dim=2,
noise_dim=2,
h_dim=64,
z_dim=2,
nonlinearity='tanh',
num_hidden_layers=1,
enc_input=False,
enc_noise=False,
clip_logvar=None,
):
super().__init__(
input_dim=input_dim,
noise_dim=noise_dim,
h_dim=h_dim,
z_dim=z_dim,
nonlinearity=nonlinearity,
num_hidden_layers=num_hidden_layers,
enc_input=enc_input,
enc_noise=enc_noise,
clip_logvar=clip_logvar,
)
inp_dim = input_dim if not enc_input else h_dim
ctx_dim = noise_dim if not enc_noise else h_dim
self.inp_encode = Identity() if not enc_input \
else MLP(input_dim=input_dim, hidden_dim=h_dim, output_dim=h_dim, nonlinearity=nonlinearity, num_hidden_layers=num_hidden_layers-1, use_nonlinearity_output=True)
self.nos_encode = Identity() if not enc_noise \
else MLP(input_dim=noise_dim, hidden_dim=h_dim, output_dim=h_dim, nonlinearity=nonlinearity, num_hidden_layers=num_hidden_layers-1, use_nonlinearity_output=True)
self.fc = MLP(input_dim=inp_dim + ctx_dim,
hidden_dim=h_dim,
output_dim=h_dim,
nonlinearity=nonlinearity,
num_hidden_layers=num_hidden_layers - 1,
use_nonlinearity_output=True)
self.reparam = NormalDistributionLinear(h_dim,
z_dim,
nonlinearity=clip_logvar)
def __init__(
self,
input_dim=2,
z_dim=2,
noise_dim=2,
h_dim=64,
nonlinearity='tanh',
num_hidden_layers=1,
enc_input=False,
enc_latent=False,
clip_logvar=None,
):
super().__init__()
self.input_dim = input_dim
self.z_dim = z_dim
self.noise_dim = noise_dim
self.h_dim = h_dim
self.nonlinearity = nonlinearity
self.num_hidden_layers = num_hidden_layers
self.enc_input = enc_input
self.enc_latent = enc_latent
inp_dim = input_dim if not enc_input else h_dim
ltt_dim = z_dim if not enc_latent else h_dim
self.inp_encode = Identity() if not enc_input \
else MLP(input_dim=input_dim, hidden_dim=h_dim, output_dim=h_dim, nonlinearity=nonlinearity, num_hidden_layers=num_hidden_layers-1, use_nonlinearity_output=True)
self.ltt_encode = Identity() if not enc_latent \
else MLP(input_dim=z_dim, hidden_dim=h_dim, output_dim=h_dim, nonlinearity=nonlinearity, num_hidden_layers=num_hidden_layers-1, use_nonlinearity_output=True)
self.fc = MLP(input_dim=inp_dim + ltt_dim,
hidden_dim=h_dim,
output_dim=h_dim,
nonlinearity=nonlinearity,
num_hidden_layers=num_hidden_layers - 1,
use_nonlinearity_output=True)
self.reparam = NormalDistributionLinear(h_dim,
noise_dim,
nonlinearity=clip_logvar)
def __init__(self,
in_dim,
out_dim,
aggregators,
scalers,
avg_d,
dropout,
batch_norm,
residual,
posttrans_layers=1):
"""
A simpler version of PNA layer that simply aggregates the neighbourhood (similar to GCN and GIN),
without using the pretransformation or the tower mechanisms of the MPNN. It does not support edge features.
:param in_dim: size of the input per node
:param out_dim: size of the output per node
:param aggregators: set of aggregation function identifiers
:param scalers: set of scaling functions identifiers
:param avg_d: average degree of nodes in the training set, used by scalers to normalize
:param dropout: dropout used
:param batch_norm: whether to use batch normalisation
:param posttrans_layers: number of layers in the transformation after the aggregation
"""
super().__init__()
# retrieve the aggregators and scalers functions
aggregators = [AGGREGATORS[aggr] for aggr in aggregators.split()]
scalers = [SCALERS[scale] for scale in scalers.split()]
self.aggregators = aggregators
self.scalers = scalers
self.in_dim = in_dim
self.out_dim = out_dim
self.dropout = dropout
self.batch_norm = batch_norm
self.residual = residual
self.batchnorm_h = nn.BatchNorm1d(out_dim)
self.posttrans = MLP(in_size=(len(aggregators) * len(scalers)) *
in_dim,
hidden_size=out_dim,
out_size=out_dim,
layers=posttrans_layers,
mid_activation='relu',
last_activation='none',
dropout=dropout)
self.avg_d = avg_d
def __init__(self, in_features, out_features, fc_layers=2, device='cpu'):
"""
:param in_features: size of the input per node
:param out_features: size of the output per node
:param fc_layers: number of fully connected layers after the sum aggregator
:param device: device used for computation
"""
super(GINLayer, self).__init__()
self.device = device
self.in_features = in_features
self.out_features = out_features
self.epsilon = nn.Parameter(torch.zeros(size=(1,), device=device))
self.post_transformation = MLP(in_size=in_features, hidden_size=max(in_features, out_features),
out_size=out_features, layers=fc_layers, mid_activation='relu',
last_activation='relu', mid_b_norm=True, last_b_norm=False, device=device)
self.reset_parameters()
def __init__(
self,
input_height=28,
input_channels=1,
z_dim=32,
nonlinearity='softplus',
#do_trim=True,
):
super().__init__()
self.input_height = input_height
self.input_channels = input_channels
self.z_dim = z_dim
self.nonlinearity = nonlinearity
#self.do_trim = do_trim
s_h = input_height
s_h2 = conv_out_size(s_h, 5, 2, 2)
s_h4 = conv_out_size(s_h2, 5, 2, 2)
s_h8 = conv_out_size(s_h4, 5, 2, 2)
#print(s_h, s_h2, s_h4, s_h8)
#_s_h8 = s_h8
#_s_h4 = deconv_out_size(_s_h8, 5, 2, 2, 0)
#_s_h2 = deconv_out_size(_s_h4+1, 5, 2, 2, 0)
#_s_h = deconv_out_size(_s_h2, 5, 2, 2, 0)
#if self.do_trim:
#else:
# _s_h = deconv_out_size(_s_h2, 5, 2, 2, 1)
#print(_s_h, _s_h2, _s_h4, _s_h8)
#ipdb.set_trace()
self.s_h8 = s_h8
self.afun = get_nonlinear_func(nonlinearity)
self.fc = MLP(input_dim=z_dim,
hidden_dim=300,
output_dim=s_h8 * s_h8 * 32,
nonlinearity=nonlinearity,
num_hidden_layers=1,
use_nonlinearity_output=True)
self.deconv1 = nn.ConvTranspose2d(32, 32, 5, 2, 2, 0, bias=True)
self.pad1 = nn.ZeroPad2d((0, 1, 0, 1))
self.deconv2 = nn.ConvTranspose2d(32, 16, 5, 2, 2, 0, bias=True)
self.reparam = BernoulliDistributionConvTranspose2d(
16, self.input_channels, 5, 2, 2, 0, bias=True)
self.padr = nn.ZeroPad2d((0, -1, 0, -1))
def __init__(
self,
input_dim=784,
h_dim=300,
z_dim=32,
nonlinearity='softplus',
num_hidden_layers=1,
):
super().__init__()
self.input_dim = input_dim
self.h_dim = h_dim
self.z_dim = z_dim
self.nonlinearity = nonlinearity
self.num_hidden_layers = num_hidden_layers
self.main = MLP(input_dim=z_dim,
hidden_dim=h_dim,
output_dim=h_dim,
nonlinearity=nonlinearity,
num_hidden_layers=num_hidden_layers,
use_nonlinearity_output=True)
self.reparam = BernoulliDistributionLinear(h_dim, input_dim)
def __init__(
self,
input_dim=2,
h_dim=1000,
std=0.1,
num_hidden_layers=1,
nonlinearity='tanh',
noise_type='gaussian',
):
super().__init__()
self.input_dim = input_dim
self.h_dim = h_dim
self.std = std
self.num_hidden_layers = num_hidden_layers
self.nonlinearity = nonlinearity
self.noise_type = noise_type
self.main = MLP(input_dim,
h_dim,
input_dim,
use_nonlinearity_output=False,
num_hidden_layers=num_hidden_layers,
nonlinearity=nonlinearity)
def __init__(
self,
noise_dim=100,
z_dim=32,
c_dim=512, #450,
h_dim=800,
num_hidden_layers=1,
nonlinearity='elu', #act=nn.ELU(),
do_center=False,
enc_noise=False,
enc_type='mlp',
):
super().__init__()
self.noise_dim = noise_dim
self.z_dim = z_dim
self.c_dim = c_dim
self.h_dim = h_dim
self.num_hidden_layers = num_hidden_layers
assert num_hidden_layers > 0
self.nonlinearity = nonlinearity
self.do_center = do_center
self.enc_noise = enc_noise
nos_dim = noise_dim if not enc_noise else c_dim
self.enc_type = enc_type
assert enc_type in [
'mlp', 'res-wn-mlp', 'res-mlp', 'res-wn-mlp-lin', 'res-mlp-lin'
]
act = get_afunc(nonlinearity_type=nonlinearity)
self.inp_encode = nn.Sequential(
nn_.ResConv2d(1, 16, 3, 2, padding=1, activation=act), act,
nn_.ResConv2d(16, 16, 3, 1, padding=1, activation=act), act,
nn_.ResConv2d(16, 32, 3, 2, padding=1, activation=act), act,
nn_.ResConv2d(32, 32, 3, 1, padding=1, activation=act), act,
nn_.ResConv2d(32, 32, 3, 2, padding=1, activation=act), act,
nn_.Reshape((-1, 32 * 4 * 4)), nn_.ResLinear(32 * 4 * 4, c_dim),
act)
#self.fc = nn.Sequential(
# nn.Linear(c_dim + nos_dim, h_dim, bias=True),
# act,
# nn.Linear(h_dim, z_dim, bias=True),
# )
if enc_type == 'mlp':
self.fc = MLP(input_dim=c_dim + nos_dim,
hidden_dim=h_dim,
output_dim=z_dim,
nonlinearity=nonlinearity,
num_hidden_layers=num_hidden_layers,
use_nonlinearity_output=False)
elif enc_type == 'res-wn-mlp':
self.fc = ResMLP(input_dim=c_dim + nos_dim,
hidden_dim=h_dim,
output_dim=z_dim,
nonlinearity=nonlinearity,
num_hidden_layers=num_hidden_layers,
use_nonlinearity_output=False,
layer='wnlinear')
elif enc_type == 'res-mlp':
self.fc = ResMLP(input_dim=c_dim + nos_dim,
hidden_dim=h_dim,
output_dim=z_dim,
nonlinearity=nonlinearity,
num_hidden_layers=num_hidden_layers,
use_nonlinearity_output=False,
layer='linear')
elif enc_type == 'res-wn-mlp-lin':
self.fc = nn.Sequential(
ResMLP(input_dim=c_dim + nos_dim,
hidden_dim=h_dim,
output_dim=h_dim,
nonlinearity=nonlinearity,
num_hidden_layers=num_hidden_layers - 1,
use_nonlinearity_output=True,
layer='wnlinear'),
nn.Linear(h_dim, z_dim, bias=True),
)
elif enc_type == 'res-mlp-lin':
self.fc = nn.Sequential(
ResMLP(input_dim=c_dim + nos_dim,
hidden_dim=h_dim,
output_dim=h_dim,
nonlinearity=nonlinearity,
num_hidden_layers=num_hidden_layers - 1,
use_nonlinearity_output=True,
layer='linear'),
nn.Linear(h_dim, z_dim, bias=True),
)
else:
raise NotImplementedError
self.nos_encode = Identity() if not enc_noise \
else nn.Sequential(
nn.Linear(noise_dim, c_dim, bias=True),
act,
)
def __init__(self,
num_inputs,
num_actions,
hidden_dim,
noise_dim,
num_enc_layers,
num_fc_layers,
args,
nonlinearity='elu',
fc_type='mlp'):
super(StochasticPolicy, self).__init__()
self.num_enc_layers = num_enc_layers
self.num_fc_layers = num_fc_layers
self.args = args
self.num_actions = num_actions
self.noise_dim = noise_dim
assert noise_dim >= num_actions, 'noise_dim: {}, num_actions: {}'.format(
noise_dim, num_actions)
inp_dim = num_inputs if num_enc_layers < 0 else hidden_dim
if num_enc_layers < 0:
self.encode = Identity()
else:
self.encode = MLP(num_inputs,
hidden_dim,
hidden_dim,
nonlinearity=nonlinearity,
num_hidden_layers=num_enc_layers,
use_nonlinearity_output=True)
if fc_type == 'mlp':
self.fc = MLP(inp_dim + noise_dim,
hidden_dim,
num_actions,
nonlinearity=nonlinearity,
num_hidden_layers=num_fc_layers,
use_nonlinearity_output=False)
torch.nn.init.normal_(self.fc.fc.weight, std=1.)
elif fc_type == 'wnres':
self.fc = ResMLP(inp_dim + noise_dim,
hidden_dim,
num_actions,
nonlinearity=nonlinearity,
num_hidden_layers=num_fc_layers,
layer='wnlinear',
use_nonlinearity_output=False)
elif fc_type == 'res':
self.fc = ResMLP(inp_dim + noise_dim,
hidden_dim,
num_actions,
nonlinearity=nonlinearity,
num_hidden_layers=num_fc_layers,
layer='linear',
use_nonlinearity_output=False)
elif fc_type == 'mlpdeep':
self.fc = MLP(inp_dim + noise_dim,
64,
num_actions,
nonlinearity=nonlinearity,
num_hidden_layers=num_fc_layers,
use_nonlinearity_output=False)
torch.nn.init.normal_(self.fc.fc.weight, std=1.)
elif fc_type == 'wnresdeep':
self.fc = ResMLP(inp_dim + noise_dim,
64,
num_actions,
nonlinearity=nonlinearity,
num_hidden_layers=num_fc_layers,
layer='wnlinear',
use_nonlinearity_output=False)
elif fc_type == 'resdeep':
self.fc = ResMLP(inp_dim + noise_dim,
64,
num_actions,
nonlinearity=nonlinearity,
num_hidden_layers=num_fc_layers,
layer='linear',
use_nonlinearity_output=False)
else:
raise NotImplementedError
def __init__(self,
in_channels,
out_channels,
aggregators,
scalers,
avg_d,
towers=1,
pretrans_layers=1,
posttrans_layers=1,
divide_input=False,
**kwargs):
"""
:param in_channels: size of the input per node
:param in_channels: size of the output per node
:param aggregators: set of aggregation function identifiers
:param scalers: set of scaling functions identifiers
:param avg_d: average degree of nodes in the training set, used by scalers to normalize
:param towers: number of towers to use
:param pretrans_layers: number of layers in the transformation before the aggregation
:param posttrans_layers: number of layers in the transformation after the aggregation
:param divide_input: whether the input features should be split between towers or not
"""
super(PNAConv, self).__init__(aggr=None, **kwargs)
assert (
(not divide_input) or in_channels % towers == 0
), "if divide_input is set the number of towers has to divide in_features"
assert (
out_channels %
towers == 0), "the number of towers has to divide the out_features"
self.in_channels = in_channels
self.out_channels = out_channels
self.towers = towers
self.divide_input = divide_input
self.input_tower = self.in_channels // towers if divide_input else self.in_channels
self.output_tower = self.out_channels // towers
# retrieve the aggregators and scalers functions
self.aggregators = [AGGREGATORS[aggr] for aggr in aggregators]
self.scalers = [SCALERS[scale] for scale in scalers]
self.avg_d = avg_d
self.edge_encoder = FCLayer(in_size=in_channels,
out_size=self.input_tower,
activation='none')
# build pre-transformations and post-transformation MLP for each tower
self.pretrans = nn.ModuleList()
self.posttrans = nn.ModuleList()
for _ in range(towers):
self.pretrans.append(
MLP(in_size=3 * self.input_tower,
hidden_size=self.input_tower,
out_size=self.input_tower,
layers=pretrans_layers,
mid_activation='relu',
last_activation='none'))
self.posttrans.append(
MLP(in_size=(len(self.aggregators) * len(self.scalers) + 1) *
self.input_tower,
hidden_size=self.output_tower,
out_size=self.output_tower,
layers=posttrans_layers,
mid_activation='relu',
last_activation='none'))
self.mixing_network = FCLayer(self.out_channels,
self.out_channels,
activation='LeakyReLU')