def __init__(self,
statistics_group_size=1,
name=None,
data_layout='data_parallel'):
super().__init__()
BatchNormModule.global_count += 1
self.instance = 0
self.statistics_group_size = statistics_group_size
self.name = (name if name else 'bnmodule{0}'.format(
BatchNormModule.global_count))
self.data_layout = data_layout
# Initialize weights
self.scale = lbann.Weights(
initializer=lbann.ConstantInitializer(value=1.0),
name=self.name + '_scale')
self.bias = lbann.Weights(
initializer=lbann.ConstantInitializer(value=0.0),
name=self.name + '_bias')
self.running_mean = lbann.Weights(
initializer=lbann.ConstantInitializer(value=0.0),
name=self.name + '_running_mean')
self.running_variance = lbann.Weights(
initializer=lbann.ConstantInitializer(value=1.0),
name=self.name + '_running_variance')
def __init__(self,
size,
statistics_group_size=1,
name=None,
data_layout='data_parallel'):
super().__init__()
FcBnRelu.global_count += 1
self.instance = 0
self.name = (name
if name else 'fcbnrelu{0}'.format(FcBnRelu.global_count))
self.data_layout = data_layout
self.fc = lbann.modules.FullyConnectedModule(
size,
bias=False,
name=self.name + '_fc',
data_layout=self.data_layout)
# Weights for batchnorm
scalebias_vals = [1.0] * size + [0.0] * size
self.bn_weights = [
lbann.Weights(name='{0}_bn_running_mean'.format(self.name),
initializer=lbann.ConstantInitializer(value=0.0)),
lbann.Weights(name='{0}_bn_running_var'.format(self.name),
initializer=lbann.ConstantInitializer(value=1.0)),
lbann.Weights(name='{0}_bn_scalebias'.format(self.name),
initializer=lbann.ValueInitializer(
values=' '.join([str(x)
for x in scalebias_vals])))
]
def __init__(self, embed_dim, num_heads, branches, d_kv=None, name=None):
super().__init__()
MultiheadAttention.global_count += 1
self.instance = 0
assert embed_dim % num_heads == 0, 'embed_dim must be divisible by num_heads'
self.embed_dim = embed_dim
self.num_heads = num_heads
self.head_dim = embed_dim // num_heads
if (d_kv == None):
self.inner_dim = embed_dim
self.head_dim = embed_dim // num_heads
else:
self.inner_dim = d_kv * num_heads
self.head_dim = d_kv
if (branches == 0):
self.ENABLE_SUBGRAPH = False
self.BRANCHES = 0
else:
self.ENABLE_SUBGRAPH = True
self.BRANCHES = branches
# Module name
self.name = name
if not self.name:
self.name = f'multiheadattention{MultiheadAttention.global_count}'
# Weights for fully-connected layers
self.query_weights = [
lbann.Weights(initializer=lbann.GlorotNormalInitializer(),
name=f'{self.name}_query_matrix'),
lbann.Weights(initializer=lbann.ConstantInitializer(value=0),
name=f'{self.name}_query_bias'),
]
self.key_weights = [
lbann.Weights(initializer=lbann.GlorotNormalInitializer(),
name=f'{self.name}_key_matrix'),
lbann.Weights(initializer=lbann.ConstantInitializer(value=0),
name=f'{self.name}_key_bias'),
]
self.value_weights = [
lbann.Weights(initializer=lbann.GlorotNormalInitializer(),
name=f'{self.name}_value_matrix'),
lbann.Weights(initializer=lbann.ConstantInitializer(value=0),
name=f'{self.name}_value_bias'),
]
#Channelwise FC in SubGraph
self.output_weights = []
for head in range(branches):
self.output_weights.append([
lbann.Weights(initializer=lbann.GlorotNormalInitializer(),
name=f'{self.name}_head{head}_output_matrix'),
lbann.Weights(initializer=lbann.ConstantInitializer(value=0),
name=f'{self.name}_head{head}_output_bias'),
])
def __init__(self, out_channels, kernel_size, stride, padding, use_bn,
bn_zero_init, bn_statistics_group_size, activation,
parallel_strategy, name, conv_weights):
"""Initialize ConvBNRelu module.
Args:
out_channels (int): Number of output channels, i.e. number
of convolution filters.
kernel_size (int): Size of convolution kernel.
stride (int): Convolution stride.
padding (int): Convolution padding.
use_bn (bool): Whether or not batch normalization layers are used.
bn_zero_init (bool): Zero-initialize batch normalization
scale.
bn_statistics_group_size (int): Aggregation size for batch
normalization statistics.
activation (lbann.Layer): The activation function.
name (str): Module name.
conv_weights (lbann.Weights): Pre-defined weights.
"""
super().__init__()
self.name = name
self.instance = 0
self.stride = stride
self.bn_statistics_group_size = bn_statistics_group_size
self.activation = activation
self.use_bn = use_bn
self.conv_weights = conv_weights
self.ps = parallel_strategy
# Initialize convolution
self.conv = lbann.modules.Convolution3dModule(
out_channels,
kernel_size,
stride=1,
padding=padding,
bias=False,
parallel_strategy=self.ps,
weights=self.conv_weights,
name=self.name + '_conv')
# Initialize batch normalization
if self.use_bn:
bn_scale_init = 0.0 if bn_zero_init else 1.0
bn_scale = lbann.Weights(
initializer=lbann.ConstantInitializer(value=bn_scale_init),
name=self.name + '_bn_scale')
bn_bias = lbann.Weights(
initializer=lbann.ConstantInitializer(value=0.0),
name=self.name + '_bn_bias')
self.bn_weights = [bn_scale, bn_bias]
def __init__(self, *args, **kwargs):
super().__init__()
self.name = kwargs["name"]
self.activation = None if "activation" not in kwargs.keys() \
else kwargs["activation"]
kwargs["activation"] = None
self.conv = lm.Convolution3dModule(*args, **kwargs)
bn_scale = lbann.Weights(
initializer=lbann.ConstantInitializer(value=1.0),
name="{}_bn_scale".format(self.name))
bn_bias = lbann.Weights(
initializer=lbann.ConstantInitializer(value=0.0),
name="{}_bn_bias".format(self.name))
self.bn_weights = [bn_scale, bn_bias]
self.instance = 0
def forward_discriminator2(self,img):
'''
Discriminator 2. Weights are frozen as part of Adversarial network = Stacked G + D
'''
bn_wts=[lbann.Weights(initializer=lbann.ConstantInitializer(value=1.0)),
lbann.Weights(initializer=lbann.ConstantInitializer(value=0.0))]
for count,lyr in enumerate(self.d2_conv):
if count==0: x=lbann.LeakyRelu(lyr(img), negative_slope=0.2)
else : x = lbann.LeakyRelu(lyr(x), negative_slope=0.2)
#### without convbrlu
# if count==0: x = lbann.LeakyRelu(lbann.BatchNormalization(lyr(img),weights=bn_wts,statistics_group_size=-1),negative_slope=0.2)
# else: x = lbann.LeakyRelu(lbann.BatchNormalization(lyr(x),weights=bn_wts,statistics_group_size=-1),negative_slope=0.2)
dims=524288
y= self.d2_fc(lbann.Reshape(x,dims=str(dims)))
return y
def __init__(self,
size,
bias=False,
weights=[],
activation=None,
transpose=False,
name=None,
parallel_strategy={}):
"""Initalize channelwise fully connected module
Args:
size (int or list): Dimension of the output tensor
bias (bool): Whether to apply bias after linearity.
transpose (bool): Whether to apply transpose of weights
matrix.
weights (`Weights` or iterator of `Weights`): Weights in
fully-connected layer. There are at most two: the
matrix and the bias. If weights are not provided, the
matrix will be initialized with He normal
initialization and the bias with zeros.
activation (type): Layer class for activation function.
name (str): Default name is in the form 'channelwisefc<index>'.
parallel_strategy (dict): Data partitioning scheme.
"""
super().__init__()
ChannelwiseFullyConnectedModule.global_count += 1
self.instance = 0
self.size = size
self.bias = bias
self.transpose = transpose
self.parallel_strategy = parallel_strategy
self.name = (name if name else 'channelwisefc{0}'.format(
ChannelwiseFullyConnectedModule.global_count))
self.data_layout = 'data_parallel'
self.weights = list(make_iterable(weights))
if len(self.weights) > 2:
raise ValueError('`FullyConnectedModule` has '
'at most two weights, '
'but got {0}'.format(len(self.weights)))
if len(self.weights) == 0:
self.weights.append(
lbann.Weights(initializer=lbann.HeNormalInitializer(),
name=self.name + '_matrix'))
if self.bias and len(self.weights) == 1:
self.weights.append(
lbann.Weights(initializer=lbann.ConstantInitializer(value=0.0),
name=self.name + '_bias'))
self.activation = None
if activation:
if isinstance(activation, type):
self.activation = activation
else:
self.activation = type(activation)
if not issubclass(self.activation, lbann.Layer):
raise ValueError('activation must be a layer')
def __init__(self, out_channels, kernel_size, stride, padding,
bn_zero_init, bn_statistics_group_size, relu, name):
"""Initialize ConvBNRelu module.
Args:
out_channels (int): Number of output channels, i.e. number
of convolution filters.
kernel_size (int): Size of convolution kernel.
stride (int): Convolution stride.
padding (int): Convolution padding.
bn_zero_init (bool): Zero-initialize batch normalization
scale.
bn_statistics_group_size (int): Group size for aggregating
batch normalization statistics.
relu (bool): Apply ReLU activation.
name (str): Module name.
"""
super().__init__()
self.name = name
self.instance = 0
# Initialize convolution
self.conv = lbann.modules.Convolution2dModule(out_channels,
kernel_size,
stride=stride,
padding=padding,
bias=False,
name=self.name + '_conv')
# Initialize batch normalization
bn_scale_init = 0.0 if bn_zero_init else 1.0
bn_scale = lbann.Weights(
initializer=lbann.ConstantInitializer(value=bn_scale_init),
name=self.name + '_bn_scale')
bn_bias = lbann.Weights(
initializer=lbann.ConstantInitializer(value=0.0),
name=self.name + '_bn_bias')
self.bn_weights = [bn_scale, bn_bias]
self.bn_statistics_group_size = bn_statistics_group_size
# Initialize ReLU
self.relu = relu
def __init__(self, size, bias=True, weights=[], activation=None,
name=None, data_layout='data_parallel'):
"""Initialize fully-connected module.
Args:
size (int): Size of output tensor.
activation (type): Layer class for activation function.
bias (bool): Whether to apply bias after linearity.
weights (`Weights` or iterator of `Weights`): Weights in
fully-connected layer. There are at most two: the
matrix and the bias. If weights are not provided, the
matrix will be initialized with He normal
initialization and the bias with zeros.
name (str): Default name is in the form 'fcmodule<index>'.
data_layout (str): Data layout.
"""
super().__init__()
FullyConnectedModule.global_count += 1
self.instance = 0
self.size = size
self.bias = bias
self.name = (name
if name
else 'fcmodule{0}'.format(FullyConnectedModule.global_count))
self.data_layout = data_layout
# Initialize weights
# Note: If weights are not provided, matrix weights are
# initialized with He normal scheme and bias weights are
# initialized with zeros.
self.weights = list(make_iterable(weights))
if len(self.weights) > 2:
raise ValueError('`FullyConnectedModule` has '
'at most two weights, '
'but got {0}'.format(len(self.weights)))
if len(self.weights) == 0:
self.weights.append(
lbann.Weights(initializer=lbann.HeNormalInitializer(),
name=self.name+'_matrix'))
if len(self.weights) == 1:
self.weights.append(
lbann.Weights(initializer=lbann.ConstantInitializer(value=0.0),
name=self.name+'_bias'))
# Initialize activation layer
self.activation = None
if activation:
if isinstance(activation, type):
self.activation = activation
else:
self.activation = type(activation)
if not issubclass(self.activation, lbann.Layer):
raise ValueError('activation must be a layer')
def forward_discriminator1(self, img):
'''
Discriminator 1
'''
bn_wts = [
lbann.Weights(initializer=lbann.ConstantInitializer(value=1.0)),
lbann.Weights(initializer=lbann.ConstantInitializer(value=0.0))
]
for count, lyr in enumerate(self.d1_conv):
if count == 0: x = lbann.LeakyRelu(lyr(img), negative_slope=0.2)
else: x = lbann.LeakyRelu(lyr(x), negative_slope=0.2)
#### without convbrlu
# if count==0: x = lbann.LeakyRelu(lbann.BatchNormalization(lyr(img),weights=bn_wts,statistics_group_size=-1),negative_slope=0.2)
# else: x = lbann.LeakyRelu(lbann.BatchNormalization(lyr(x),weights=bn_wts,statistics_group_size=-1),negative_slope=0.2)
dims = 32768
#dims=25088 ## for padding=1
y = self.d1_fc(lbann.Reshape(x, dims=str(dims)))
return y
def __init__(self, embed_dim, num_heads, name=None):
super().__init__()
MultiheadAttention.global_count += 1
self.instance = 0
assert embed_dim % num_heads == 0, 'embed_dim must be divisible by num_heads'
self.embed_dim = embed_dim
self.num_heads = num_heads
self.head_dim = embed_dim // num_heads
# Module name
self.name = name
if not self.name:
self.name = f'multiheadattention{MultiheadAttention.global_count}'
# Weights for fully-connected layers
self.query_weights = [
lbann.Weights(initializer=lbann.GlorotNormalInitializer(),
name=f'{self.name}_query_matrix'),
lbann.Weights(initializer=lbann.ConstantInitializer(value=0),
name=f'{self.name}_query_bias'),
]
self.key_weights = [
lbann.Weights(initializer=lbann.GlorotNormalInitializer(),
name=f'{self.name}_key_matrix'),
lbann.Weights(initializer=lbann.ConstantInitializer(value=0),
name=f'{self.name}_key_bias'),
]
self.value_weights = [
lbann.Weights(initializer=lbann.GlorotNormalInitializer(),
name=f'{self.name}_value_matrix'),
lbann.Weights(initializer=lbann.ConstantInitializer(value=0),
name=f'{self.name}_value_bias'),
]
self.output_weights = [
lbann.Weights(initializer=lbann.GlorotNormalInitializer(),
name=f'{self.name}_output_matrix'),
lbann.Weights(initializer=lbann.ConstantInitializer(value=0),
name=f'{self.name}_output_bias'),
]
def __init__(
self,
embed_dim=512,
num_heads=8,
feedforward_dim=2048,
dropout=0.1,
name=None,
):
TransformerDecoderLayer.global_count += 1
self.instance = 0
self.embed_dim = embed_dim
self.feedforward_dim = feedforward_dim
self.dropout_prob = dropout
# Module name
self.name = name
if not self.name:
self.name = f'transformerdecoderlayer{TransformerDecoderLayer.global_count}'
# Layer modules
self.attention1 = lbann.modules.transformer.MultiheadAttention(
embed_dim, num_heads, name=f'{self.name}_attention1')
self.attention2 = lbann.modules.transformer.MultiheadAttention(
embed_dim, num_heads, name=f'{self.name}_attention2')
# Weights for fully-connected layers
self.fc1_weights = [
lbann.Weights(initializer=lbann.HeNormalInitializer(),
name=f'{self.name}_fc1_matrix'),
lbann.Weights(initializer=lbann.ConstantInitializer(value=0),
name=f'{self.name}_fc1_bias'),
]
self.fc2_weights = [
lbann.Weights(initializer=lbann.GlorotNormalInitializer(),
name=f'{self.name}_fc2_matrix'),
lbann.Weights(initializer=lbann.ConstantInitializer(value=0),
name=f'{self.name}_fc2_bias'),
]
def __init__(self,
neuron_dims=[1000, 1000, 1000],
activation=lbann.Relu,
keep_prob=0.95,
name=None):
self.instance = 0
self.name = (name if name else 'combo{0}'.format(Combo.global_count))
#shared weights for drug 1 and 2 tracks
shared_w = []
for i in range(len(neuron_dims)):
shared_w.append(
lbann.Weights(initializer=lbann.HeNormalInitializer(),
name='drug_matrix' + str(i)))
shared_w.append(
lbann.Weights(initializer=lbann.ConstantInitializer(value=0.0),
name='drug_bias' + str(i)))
print("SHARED W ", type(shared_w))
self.geneT = TrackModule(neuron_dims,
activation,
keep_prob,
name=self.name + 'gene_track')
self.drug1T = TrackModule(neuron_dims,
activation,
keep_prob,
shared_w,
name=self.name + 'drug1_track')
self.drug2T = TrackModule(neuron_dims,
activation,
keep_prob,
shared_w,
name=self.name + 'drug2_track')
self.concatT = TrackModule(neuron_dims,
activation,
keep_prob,
name=self.name + 'concat_track')
def __init__(self,
input_channels,
output_channels,
num_nodes,
bias=True,
activation=lbann.Relu,
name=None,
parallel_strategy={}):
"""Initialize GCN layer
Args:
input_channels (int): The size of the input node features
output_channels (int): The output size of the node features
num_nodes (int): Number of vertices in the graph
bias (bool): Whether to apply biases after weights transform
activation (type): Activation leyer for the node features. If None, then no activation is
applied. (default: lbann.Relu)
name (str): Default name of the layer is GCN_{number}
parallel_strategy (dict): Data partitioning scheme.
"""
super().__init__()
## Add variables
self.input_channel_size = input_channels
self.output_channel_size = output_channels
self.num_nodes = num_nodes
self.parallel_strategy = parallel_strategy
self.instance = 0
self.is_distconv = False
if parallel_strategy:
if list(parallel_strategy.values()[0]) > 0:
self.is_distconv = True
## Add Name for the components for the layer
GCNConv.global_count += 1
self.name = (name if name else 'GCN_{}'.format(GCNConv.global_count))
weights = []
## Initialize weights for the matrix
value = math.sqrt(6 / (input_channels + output_channels))
weights.append(
lbann.Weights(initializer=lbann.UniformInitializer(min=-value,
max=value),
name=self.name + '_weights'))
## Initialize bias variables
self.has_bias = bias
if (self.has_bias):
weights.append(
lbann.Weights(initializer=lbann.ConstantInitializer(value=0.0),
name=self.name + '_bias_weights'))
self.activation = None
if activation:
if isinstance(activation, type):
self.activation = activation
else:
self.activation = type(actvation)
if not issubclass(self.activation, lbann.Layer):
raise ValueError('activation must be a layer')
# Distconv channelwise fully connected expects 3D tensors as input
# and output. This check adds an extra dimention to enable
# channel-wise data partitioning
self.output_channels = self.output_channel_size
if self.is_distconv:
self.output_channels = [1, self.output_channel_size]
self.nn = \
ChannelwiseFullyConnectedModule(self.output_channels,
bias=self.has_bias,
weights=weights,
activation=self.activation,
name=self.name+"_FC_layer",
parallel_strategy=self.parallel_strategy)
def __init__(self,
input_width,
output_size,
name=None,
use_bn=False,
bn_statistics_group_size=None):
"""Initialize CosmFlow.
Args:
input_width (int): Size of each spatial dimension of input data.
output_size (int): Size of output tensor.
name (str, optional): Module name
(default: 'cosmoflow_module<index>').
use_bn (bool): Whether or not batch normalization layers are used.
bn_statistics_group_size (int): The number of samples
for each batch-normalization group.
"""
CosmoFlow.global_count += 1
self.instance = 0
self.name = (name if name else 'cosmoflow_module{0}'.format(
CosmoFlow.global_count))
self.input_width = input_width
self.use_bn = use_bn
assert self.input_width in [128, 256, 512]
self.cp_params = [
{
"type": "conv",
"out_channels": 16,
"kernel_size": 3,
"stride": 1
},
{
"type": "pool"
},
{
"type": "conv",
"out_channels": 32,
"kernel_size": 3,
"stride": 1
},
{
"type": "pool"
},
{
"type": "conv",
"out_channels": 64,
"kernel_size": 3,
"stride": 1
},
{
"type": "pool"
},
{
"type": "conv",
"out_channels": 128,
"kernel_size": 3,
"stride": 2
},
{
"type": "pool"
},
{
"type": "conv",
"out_channels": 256,
"kernel_size": 3,
"stride": 1
},
{
"type": "pool"
},
{
"type": "conv",
"out_channels": 256,
"kernel_size": 3,
"stride": 1
},
{
"type": "conv",
"out_channels": 256,
"kernel_size": 3,
"stride": 1
},
]
additional_pools = []
if self.input_width == 256:
additional_pools = [6]
elif self.input_width == 512:
additional_pools = [6, 7]
for i in additional_pools:
conv_idx = list(
np.cumsum([
1 if x["type"] == "conv" else 0 for x in self.cp_params
])).index(i)
self.cp_params.insert(conv_idx + 1, {"type": "pool"})
for p in self.cp_params:
if p["type"] == "conv":
p["padding"] = int((p["kernel_size"] - 1) / 2)
# Create convolutional blocks
activation = lbann.LeakyRelu
for i, param in enumerate(
filter(lambda x: x["type"] == "conv", self.cp_params)):
conv_name = "conv" + str(i + 1)
conv_weights = [
lbann.Weights(initializer=lbann.GlorotUniformInitializer())
]
param_actual = dict(param)
param_actual.pop("type", None)
conv = ConvBNRelu(
**param_actual,
conv_weights=conv_weights,
use_bn=self.use_bn,
bn_statistics_group_size=bn_statistics_group_size,
bn_zero_init=False,
name=self.name + "_" + conv_name,
activation=lbann.LeakyRelu)
setattr(self, conv_name, conv)
# Create fully-connected layers
fc_params = [
{
"size": 2048
},
{
"size": 256
},
{
"size": output_size
},
]
for i, param in enumerate(fc_params):
fc_name = "fc" + str(i + 1)
fc = lm.FullyConnectedModule(
**param,
activation=activation if i < len(fc_params) - 1 else None,
name=self.name + "_" + fc_name,
weights=[
lbann.Weights(
initializer=lbann.GlorotUniformInitializer()),
lbann.Weights(initializer=lbann.ConstantInitializer(
value=0.1))
],
)
setattr(self, fc_name, fc)
def __init__(self,
input_channels,
output_channels,
bias=True,
activation=lbann.Relu,
name=None,
data_layout='data_parallel'):
"""Initialize Graph layer
Args:
input_channels (int): The size of the input node features
output_channels (int): The output size of the node features
bias (bool): Whether to apply biases after MatMul
name (str): Default name of the layer is GCN_{number}
data_layout (str): Data layout
activation (type): Activation layer for the node features. If None, then no activation is
applied. (default: lbann.Relu)
"""
super().__init__()
## Add variables
self.input_channels = input_channels
self.output_channels = output_channels
self.data_layout = data_layout
## Add Name for the components for the layer
GraphConv.global_count += 1
self.name = (name
if name else 'Graph_{}'.format(GraphConv.global_count))
## Initialize weights for the matrix
value = math.sqrt(6 / (input_channels + output_channels))
self.mat_weights = lbann.Weights(initializer=lbann.UniformInitializer(
min=-value, max=value),
name=self.name + '_Weights')
self.weights1 = lbann.WeightsLayer(dims=str_list(
[input_channels, output_channels]),
name=self.name + '_layer',
weights=self.mat_weights)
self.id_weights = lbann.Weights(initializer=lbann.UniformInitializer(
min=-value, max=value),
name=self.name + '_ID_Weights')
self.weights2 = lbann.WeightsLayer(dims=str_list(
[input_channels, output_channels]),
name=self.name + '_ID_layer',
weights=self.id_weights)
## Initialize bias variables
self.has_bias = bias
self.bias_weights = None
self.bias = None
if (self.has_bias):
self.bias_weights = lbann.Weights(
initializer=lbann.ConstantInitializer(value=0.0),
name=self.name + '_bias_weights')
self.bias = lbann.WeightsLayer(dims=str_list([1, output_channels]),
weights=self.bias_weights,
name=self.name + '_bias_layer')
self.activation = None
if activation:
if isinstance(activation, type):
self.activation = activation
else:
self.activation = type(actvation)
if not issubclass(self.activation, lbann.Layer):
raise ValueError('activation must be a layer')
def __init__(self,
input_channels,
output_channels,
num_nodes,
bias=True,
activation=lbann.Relu,
name=None):
"""Initialize Graph layer
Args:
input_channels (int): The size of the input node features
output_channels (int): The output size of the node features
num_nodes (int): Number of vertices in the graph
bias (bool): Whether to apply biases after weights transform
activation (type): Activation layer for the node features. If None, then no activation is
applied. (default: lbann.Relu)
name (str): Default name of the layer is Graph_{number}
"""
super().__init__()
## Add variables
self.input_channel_size = input_channels
self.output_channel_size = output_channels
self.num_nodes = num_nodes
## Add Name for the components for the layer
GraphConv.global_count += 1
self.name = (name
if name else 'Graph_{}'.format(GraphConv.global_count))
## Initialize weights for the matrix
value = math.sqrt(6 / (input_channels + output_channels))
mat_weights = []
id_weights = []
mat_weights.append(
lbann.Weights(initializer=lbann.UniformInitializer(min=-value,
max=value),
name=self.name + '_Weights'))
id_weights.append(
lbann.Weights(initializer=lbann.UniformInitializer(min=-value,
max=value),
name=self.name + '_ID_Weights'))
## Initialize bias variables
self.has_bias = bias
if (self.has_bias):
mat_weights.append(
lbann.Weights(initializer=lbann.ConstantInitializer(value=0.0),
name=self.name + '_bias_weights'))
self.activation = None
if activation:
if isinstance(activation, type):
self.activation = activation
else:
self.activation = type(actvation)
if not issubclass(self.activation, lbann.Layer):
raise ValueError('activation must be a layer')
self.id_nn = \
ChannelwiseFullyConnectedModule(self.output_channel_size,
bias=False,
weights=id_weights,
activation=self.activation,
name=self.name+"_ID_FC_layer")
self.mat_nn = \
ChannelwiseFullyConnectedModule(self.output_channel_size,
bias=self.has_bias,
weights=mat_weights,
activation=self.activation,
name=self.name+"_Message_FC_layer")
def __init__(self,
num_dims,
out_channels,
kernel_size,
stride=1,
padding=0,
dilation=1,
groups=1,
bias=True,
weights=[],
activation=None,
name=None,
transpose=False,
parallel_strategy={}):
"""Initialize convolution module.
Args:
num_dims (int): Number of dimensions.
out_channels (int): Number of output channels, i.e. number
of filters.
kernel_size (int): Size of convolution kernel.
stride (int): Convolution stride.
padding (int): Convolution padding.
dilation (int): Convolution dilation.
groups (int): Number of convolution groups.
bias (bool): Whether to apply channel-wise bias after
convolution.
weights (`Weights` or iterator of `Weights`): Weights in
convolution layer. There are at most two: the kernel
and the bias. If weights are not provided, the kernel
will be initialized with He normal initialization and
the bias with zeros.
name (str): Default name is in the form 'convmodule<index>'.
transpose (bool): If true call deconvolution (or convolution
transpose)
parallel_strategy dict): Data partitioning scheme.
"""
super().__init__()
ConvolutionModule.global_count += 1
self.instance = 0
self.num_dims = num_dims
self.out_channels = out_channels
self.kernel_size = kernel_size
self.stride = stride
self.padding = padding
self.dilation = dilation
self.groups = groups
self.bias = bias
self.weights = list(make_iterable(weights))
self.name = (name if name else 'convmodule{0}'.format(
ConvolutionModule.global_count))
self.transpose = transpose
self.parallel_strategy = parallel_strategy
# Initialize weights
# Note: If weights are not provided, kernel weights are
# initialized with He normal scheme and bias weights are
# initialized with zeros.
self.weights = list(make_iterable(weights))
if len(self.weights) > 2:
raise ValueError('`ConvolutionModule` has '
'at most two weights, '
'but got {0}'.format(len(self.weights)))
if len(self.weights) == 0:
self.weights.append(
lbann.Weights(initializer=lbann.HeNormalInitializer(),
name=self.name + '_kernel'))
if len(self.weights) == 1:
self.weights.append(
lbann.Weights(initializer=lbann.ConstantInitializer(value=0.0),
name=self.name + '_bias'))
# Initialize activation layer
self.activation = None
if activation:
if isinstance(activation, type):
self.activation = activation
else:
self.activation = type(activation)
if not issubclass(self.activation, lbann.Layer):
raise ValueError('activation must be a layer')
'--work-dir', action='store', default=None, type=str,
help='working directory', metavar='DIR')
args = parser.parse_args()
# ----------------------------------
# Embedding weights
# ----------------------------------
encoder_embeddings_weights = lbann.Weights(
initializer=lbann.NormalInitializer(
mean=0, standard_deviation=1/args.latent_dim,
),
name='embeddings',
)
decoder_embeddings_weights = lbann.Weights(
initializer=lbann.ConstantInitializer(value=0),
name='decoder_embeddings',
)
# ----------------------------------
# Construct layer graph
# ----------------------------------
# Properties of graph and random walk
num_graph_nodes = dataset.max_graph_node_id() + 1
walk_length = dataset.walk_context_length
num_negative_samples = dataset.num_negative_samples
input_size = dataset.sample_dims()[0]
# Embedding vectors, including negative sampling
# Note: Input is sequence of graph node IDs
def __init__(self, output_size, input_width, name=None):
"""Initialize CosmFlow.
Args:
output_size (int): Size of output tensor.
input_width (int): Width of input tensor.
name (str, optional): Module name
(default: 'cosmoflow_module<index>').
"""
CosmoFlow.global_count += 1
self.instance = 0
self.name = (name if name else 'cosmoflow_module{0}'.format(
CosmoFlow.global_count))
self.input_width = input_width
assert self.input_width in [128, 256, 512]
self.layer_params = [
{
"type": "conv",
"out_channels": 16,
"kernel_size": 3,
"stride": 1
},
{
"type": "pool"
},
{
"type": "conv",
"out_channels": 32,
"kernel_size": 3,
"stride": 1
},
{
"type": "pool"
},
{
"type": "conv",
"out_channels": 64,
"kernel_size": 3,
"stride": 1
},
{
"type": "pool"
},
{
"type": "conv",
"out_channels": 128,
"kernel_size": 3,
"stride": 2
},
{
"type": "pool"
},
{
"type": "conv",
"out_channels": 256,
"kernel_size": 3,
"stride": 1
},
{
"type": "pool"
},
{
"type": "conv",
"out_channels": 256,
"kernel_size": 3,
"stride": 1
},
{
"type": "conv",
"out_channels": 256,
"kernel_size": 3,
"stride": 1
},
]
for p in self.layer_params:
if p["type"] == "conv":
p["padding"] = int((p["kernel_size"] - 1) / 2)
additional_pools = []
if self.input_width == 256:
additional_pools = [6]
elif self.input_width == 512:
additional_pools = [6, 7]
for i in additional_pools:
conv_idx = list(
np.cumsum([
1 if x["type"] == "conv" else 0 for x in self.layer_params
])).index(i)
self.layer_params.insert(conv_idx + 1, {"type": "pool"})
width = self.input_width
for p in self.layer_params:
if p["type"] == "conv":
output_width = int(width / p["stride"])
else:
output_width = int(width / 2)
p["width"] = output_width
width = output_width
assert width > 0
for i, param in enumerate(
filter(lambda x: x["type"] == "conv", self.layer_params)):
conv_name = "conv" + str(i + 1)
conv_weights = [
Weights(initializer=lbann.GlorotUniformInitializer())
]
param_actual = dict(param)
param_actual.pop("type", None)
param_actual.pop("width", None)
conv = lm.Convolution3dModule(**param_actual,
activation=lbann.LeakyRelu,
name=self.name + "_" + conv_name,
bias=False,
weights=conv_weights)
setattr(self, conv_name, conv)
# Create fully-connected layers
fc_params = [
{
"size": 2048
},
{
"size": 256
},
{
"size": output_size
},
]
for i, param in enumerate(fc_params):
fc_name = "fc" + str(i + 1)
fc = lm.FullyConnectedModule(
**param,
activation=lbann.LeakyRelu if i < len(fc_params) - 1 else None,
name=self.name + "_" + fc_name,
weights=[
Weights(initializer=lbann.GlorotUniformInitializer()),
Weights(initializer=lbann.ConstantInitializer(value=0.1))
],
)
setattr(self, fc_name, fc)
