def forward_generator(self,z):
x = lbann.Relu(lbann.BatchNormalization(self.g_fc1(z),decay=0.9,scale_init=1.0,epsilon=1e-5))
x = lbann.Reshape(x, dims='512 8 8') #channel first
x = lbann.Relu(lbann.BatchNormalization(self.g_convT[0](x),decay=0.9,scale_init=1.0,epsilon=1e-5))
x = lbann.Relu(lbann.BatchNormalization(self.g_convT[1](x),decay=0.9,scale_init=1.0,epsilon=1e-5))
x = lbann.Relu(lbann.BatchNormalization(self.g_convT[2](x),decay=0.9,scale_init=1.0,epsilon=1e-5))
return self.g_convT3(x)
def forward_generator(self, z, mcr):
'''
Build the Generator
'''
x = lbann.Relu(
lbann.BatchNormalization(self.g_fc1(z),
decay=0.9,
scale_init=1.0,
epsilon=1e-5))
dims = '512 8 8'
x = lbann.Reshape(x, dims=dims) #channel first
for count, lyr in enumerate(self.g_convT):
x = lbann.Relu(
lbann.BatchNormalization(lyr(x),
decay=0.9,
scale_init=1.0,
epsilon=1e-5))
img = self.g_convT3(x)
if mcr: ### For multi-channel rescaling, add extra channel to output image
linear_scale = 1 / self.linear_scaler
# linear_scale=lbann.Constant(value=0.001)
ch2 = lbann.Tanh(
lbann.WeightedSum(self.inv_transform(img),
scaling_factors=str(linear_scale)))
y = lbann.Concatenation(img, ch2, axis=0)
img = lbann.Reshape(y, dims='2 128 128')
else:
img = lbann.Reshape(img, dims='1 128 128')
return img
def DGCN_layer(feature_matrix,adj_matrix, node_features):
"""An example 3 layer GCN kernel.
Args:
feature_matrix (Layer): Node feature layer. Should have the shape:
(num_nodes, node_features)
adj_matrix (Layer): Adjancency matrix layer. Should have the shape:
(num_nodes, num_nodes)
node_features (int): The number of features per node
Returns:
(Layer): Returns the new embedding of the node features
"""
out_channel_1 = 1024
out_channel_2 = 512
out_channel_3 = 256
gcn1 = DenseGCNConv(input_channels = node_features, output_channels = out_channel_1)
gcn2 = DenseGCNConv(input_channels = out_channel_1, output_channels = out_channel_2)
gcn3 = DenseGCNConv(input_channels = out_channel_2, output_channels = out_channel_3)
out_channel = out_channel_3
x = gcn1(feature_matrix, adj_matrix )
x = lbann.Relu(x,name="DGCN1_activation")
x = gcn2(x, adj_matrix)
x = lbann.Relu(x, name="DGCN2_activation")
x = gcn3 (x, adj_matrix)
x = lbann.Relu(x, name="DGCN3_activation")
return x
def decoder_cnn(self, z):
x = self.dec_cnn_fc(z)
sca = self.dec_fc_sca(lbann.Identity(x))
img = lbann.Reshape(lbann.Identity(x),
dims="16 8 8",
name=self.name + 'dec_reshape0')
img = self.dec_convT[2](lbann.Relu(self.dec_convT[1](lbann.Relu(
self.dec_convT[0](img)))))
#concat for common interface, slice in output
img = lbann.Reshape(img,
dims=str(64 * 64 * 4),
name=self.name +
'dec_reshape1') #?? check tensor shape
#todo check that concat size == dec_out_dim
return lbann.Concatenation([img, sca], axis=0)
def forward(self, x):
self.instance += 1
x = self.conv(x)
x = self.bn(x)
return lbann.Relu(x,
name='{0}_relu_instance{1}'.format(
self.name, self.instance))
def conv_block(statistics_group_size, current_block_num, current_layer_num,
cumulative_layer_num, parent_node, conv_dims_i, conv_pads_i,
num_output_channels):
batch_normalization_node = standard_batchnorm(statistics_group_size,
parent_node)
cumulative_layer_num += 1
log('dense_block={b} dense_layer={l} BatchNormalization. cumulative_layer_num={n}'
.format(b=current_block_num,
l=current_layer_num,
n=cumulative_layer_num))
relu_node = lbann.Relu(batch_normalization_node)
cumulative_layer_num += 1
log('dense_block={b} dense_layer={l} Relu. cumulative_layer_num={n}'.
format(b=current_block_num,
l=current_layer_num,
n=cumulative_layer_num))
convolution_node = lbann.Convolution(
relu_node,
conv_dims_i=conv_dims_i,
conv_pads_i=conv_pads_i,
conv_strides_i=1,
has_bias=False,
num_dims=2,
num_output_channels=num_output_channels)
cumulative_layer_num += 1
log('dense_block={b} dense_layer={l} Convolution. cumulative_layer_num={n}'
.format(b=current_block_num,
l=current_layer_num,
n=cumulative_layer_num))
return convolution_node, cumulative_layer_num
def forward_generator(self, z, mcr):
'''
Build the Generator
'''
x = lbann.Relu(
lbann.BatchNormalization(self.g_fc1(z),
decay=0.9,
scale_init=1.0,
epsilon=1e-5))
dims = '512 8 8'
print("dims", dims)
x = lbann.Reshape(x, dims=dims) #channel first
x = lbann.Relu(
lbann.BatchNormalization(self.g_convT[0](x),
decay=0.9,
scale_init=1.0,
epsilon=1e-5))
x = lbann.Relu(
lbann.BatchNormalization(self.g_convT[1](x),
decay=0.9,
scale_init=1.0,
epsilon=1e-5))
x = lbann.Relu(
lbann.BatchNormalization(self.g_convT[2](x),
decay=0.9,
scale_init=1.0,
epsilon=1e-5))
img = self.g_convT3(x)
if mcr: ### For multi-channel rescaling, add extra channel to output image
linear_scale = 1 / self.linear_scaler
#ch2 = lbann.Tanh(self.inv_transform(img)/linear_scalar)
ch2 = lbann.Tanh(
lbann.WeightedSum(self.inv_transform(img),
scaling_factors=str(linear_scale)))
y = lbann.Concatenation(img, ch2, axis=0)
img = lbann.Reshape(y, dims='2 128 128')
else:
img = lbann.Reshape(img, dims='1 128 128')
print('Gen Img in GAN', img.__dict__)
return img
def forward(self, x):
self.instance += 1
y1 = self.branch1(x) if self.branch1 else x
y2 = self.branch2c(self.branch2b(self.branch2a(x)))
z = lbann.Add([y1, y2],
name='{0}_sum_instance{1}'.format(
self.name, self.instance))
return lbann.Relu(z,
name='{0}_relu_instance{1}'.format(
self.name, self.instance))
def forward(self, x, mask=None):
"""Apply Transformer encoder layer.
Args:
x (lbann.Layer): Sequence of input vectors.
mask (lbann.Layer, optional): Attention mask.
Returns:
lbann.Layer: Sequence of output vectors.
"""
self.instance += 1
name = f'{self.name}_instance{self.instance}'
# Self-attention with residual connection
y = self.attention(x, x, x, mask=mask)
if self.dropout_prob > 0:
y = lbann.Dropout(
y,
keep_prob=1 - self.dropout_prob,
name=f'{name}_drop1',
)
z = lbann.Sum(x, y, name=f'{name}_sum1')
z = lbann.InstanceNorm(z, name=f'{name}_norm1')
x = z
# Feedforward network with residual connection
y = lbann.ChannelwiseFullyConnected(
x,
weights=self.fc1_weights,
output_channel_dims=[self.feedforward_dim],
name=f'{name}_fc1',
)
y = lbann.Relu(y, name=f'{name}_relu1')
if self.dropout_prob > 0:
y = lbann.Dropout(
y,
keep_prob=1 - self.dropout_prob,
name=f'{name}_drop2',
)
y = lbann.ChannelwiseFullyConnected(
y,
weights=self.fc2_weights,
output_channel_dims=[self.embed_dim],
name=f'{name}_fc2',
)
if self.dropout_prob > 0:
y = lbann.Dropout(
y,
keep_prob=1 - self.dropout_prob,
name=f'{name}_drop3',
)
z = lbann.Sum(x, y, name=f'{name}_sum2')
z = lbann.InstanceNorm(z, name=f'{name}_norm2')
return z
def test_o2l_layer_Relu(self):
N, C, H, W = (100, 200, 300, 400)
lbannRelu = lbann.Relu(lbann.Input(), )
node = onnx.helper.make_node(
"Relu",
inputs=["x"],
outputs=["y"],
)
onnxRelu = convertOnnxNode(node, {"x": [N, C, H, W]}, {}).relu
self._assertFields(lbannRelu, onnxRelu)
def forward_generator(self, z, ps=None):
x = lbann.FullyConnected(z,
num_neurons=np.prod(self.outc_dims),
has_bias=True,
device=self.g_device)
x = lbann.Reshape(x,
dims=list2str(self.outc_dims),
name='gen_zin_reshape',
device=self.g_device)
x = lbann.Relu(self.g_convT[0](x),
name='g_relu0',
parallel_strategy=ps,
device=self.g_device)
x = lbann.Relu(self.g_convT[1](x),
name='g_relu1',
parallel_strategy=ps,
device=self.g_device)
x = lbann.Relu(self.g_convT[2](x),
name='g_relu2',
parallel_strategy=ps,
device=self.g_device)
return self.g_convT3(x)
def construct_model():
"""Model description
"""
import lbann
import lbann.modules
fc = lbann.modules.FullyConnectedModule
conv = lbann.modules.Convolution2dModule
conv1 = conv(20, 3, stride=1, padding=1, name='conv1')
conv2 = conv(20, 3, stride=1, padding=1, name='conv2')
fc1 = fc(100, name='fc1')
fc2 = fc(20, name='fc2')
fc3 = fc(num_classes, name='fc3')
# Layer graph
input = lbann.Input(name='inp_tensor', target_mode='classification')
inp_slice = lbann.Slice(input,
axis=0,
slice_points=str_list([0, dims - 1, dims]),
name='inp_slice')
xdata = lbann.Identity(inp_slice)
ylabel = lbann.Identity(inp_slice, name='gt_y')
#NHWC to NCHW
x = lbann.Reshape(xdata, dims='14 13 13')
x = conv2(conv1(x))
x = lbann.Reshape(x, dims='3380')
x = lbann.Dropout(lbann.Relu(fc1(x)), keep_prob=0.5)
x = lbann.Dropout(fc2(x), keep_prob=0.5)
pred = lbann.Softmax(fc3(x))
gt_label = lbann.OneHot(ylabel, size=num_classes)
loss = lbann.CrossEntropy([pred, gt_label], name='loss')
acc = lbann.CategoricalAccuracy([pred, gt_label])
layers = list(lbann.traverse_layer_graph(input))
# Setup objective function
weights = set()
for l in layers:
weights.update(l.weights)
obj = lbann.ObjectiveFunction(loss)
callbacks = [lbann.CallbackPrint(), lbann.CallbackTimer()]
# Construct model
num_epochs = 10
return lbann.Model(num_epochs,
weights=weights,
layers=layers,
metrics=[lbann.Metric(acc, name='accuracy', unit='%')],
objective_function=obj,
callbacks=callbacks)
def forward(self, x):
self.instance += 1
conv = self.conv(x)
bn = lbann.BatchNormalization(
conv,
weights=self.bn_weights,
stats_aggregation=self.bn_stats_aggregation,
name='{0}_bn_instance{1}'.format(self.name, self.instance))
if self.relu:
return lbann.Relu(bn,
name='{0}_relu_instance{1}'.format(
self.name, self.instance))
else:
return bn
def test_l2o_layer_relu(self):
N, C, H, W = (100, 200, 300, 400)
onnxRelu = onnx.helper.make_node(
"Relu",
inputs=["x"],
outputs=["y"],
)
layer = lbann.Relu(lbann.Input(name="x"), )
lbannRelu = parseLbannLayer(layer.export_proto(),
{"x_0": (N, C, H, W)})["nodes"]
self._assertFields(lbannRelu, onnxRelu)
def forward(self, x):
self.instance += 1
conv = self.conv(x)
bn = lbann.BatchNormalization(
conv,
weights=self.bn_weights,
statistics_group_size=(-1 if self.bn_statistics_group_size == 0
else self.bn_statistics_group_size),
name='{0}_bn_instance{1}'.format(self.name, self.instance))
if self.relu:
return lbann.Relu(bn,
name='{0}_relu_instance{1}'.format(
self.name, self.instance))
else:
return bn
def encode(self, x):
x = lbann.Reshape(x, dims=utils.str_list([-1, self.input_dim]))
for i, dim in enumerate(self.hidden_dims):
x = lbann.ChannelwiseFullyConnected(
x,
weights=self.weights[i],
output_channel_dims=dim,
bias=False,
)
x = lbann.Relu(x)
x = lbann.ChannelwiseFullyConnected(
x,
weights=self.weights[-1],
output_channel_dims=self.output_dim,
bias=False,
)
return x
def decode(self, x):
x = lbann.Reshape(x, dims=utils.str_list([-1, self.output_dim]))
for i in range(len(self.hidden_dims)):
x = lbann.ChannelwiseFullyConnected(
x,
weights=self.weights[-i - 1],
output_channel_dims=self.hidden_dims[-i - 1],
transpose=True,
bias=False,
)
x = lbann.Relu(x)
x = lbann.ChannelwiseFullyConnected(
x,
weights=self.weights[0],
output_channel_dims=self.input_dim,
transpose=True,
bias=False,
)
return x
def transition_layer(current_block_num,
cumulative_layer_num,
parent_node,
num_output_channels
):
batch_normalization_node = standard_batchnorm(parent_node)
cumulative_layer_num += 1
log('dense_block={b} > transition_layer BatchNormalization. cumulative_layer_num={n}'.format(
b=current_block_num, n=cumulative_layer_num))
relu_node = lbann.Relu(batch_normalization_node)
cumulative_layer_num += 1
log('dense_block={b} > transition_layer Relu. cumulative_layer_num={n}'.format(
b=current_block_num, n=cumulative_layer_num))
convolution_node = lbann.Convolution(
relu_node,
conv_dims_i=1,
conv_pads_i=0,
conv_strides_i=1,
has_bias=False,
num_dims=2,
num_output_channels=num_output_channels
)
cumulative_layer_num += 1
log('dense_block={b} > transition_layer Convolution. cumulative_layer_num={n}'.format(
b=current_block_num, n=cumulative_layer_num))
# 2x2 average pool, stride 2
pooling_node = lbann.Pooling(
convolution_node,
num_dims=2,
pool_dims_i=2,
pool_mode='average',
pool_pads_i=0,
pool_strides_i=2
)
cumulative_layer_num += 1
log('dense_block={b} > transition_layer Pooling. cumulative_layer_num={n}'.format(
b=current_block_num, n=cumulative_layer_num))
return pooling_node, cumulative_layer_num
def initial_layer(cumulative_layer_num,
images_node,
num_initial_channels
):
# 7x7 conv, stride 2
convolution_node = lbann.Convolution(
images_node,
conv_dims_i=7,
conv_pads_i=3,
conv_strides_i=2,
has_bias=False,
num_dims=2,
num_output_channels=num_initial_channels
)
cumulative_layer_num += 1
log('initial_layer Convolution. cumulative_layer_num={n}'.format(
n=cumulative_layer_num))
batch_normalization_node = standard_batchnorm(convolution_node)
cumulative_layer_num += 1
log('initial_layer BatchNormalization. cumulative_layer_num={n}'.format(
n=cumulative_layer_num))
relu_node = lbann.Relu(batch_normalization_node)
cumulative_layer_num += 1
log('initial_layer Relu. cumulative_layer_num={n}'.format(
n=cumulative_layer_num))
# 3x3 max pool, stride 2
pooling_node = lbann.Pooling(
relu_node,
num_dims=2,
pool_dims_i=3,
pool_mode='max',
pool_pads_i=1,
pool_strides_i=2
)
cumulative_layer_num += 1
log('initial_layer Pooling. cumulative_layer_num={n}'.format(
n=cumulative_layer_num))
return pooling_node, cumulative_layer_num
def forward(self, x):
self.instance += 1
x = self.fc(x)
x = lbann.EntrywiseBatchNormalization(
x,
weights=[self.bn_weights[0], self.bn_weights[1]],
decay=0.9,
epsilon=1e-5,
name='{0}_bn_instance{1}'.format(self.name, self.instance),
data_layout=self.data_layout)
x = lbann.EntrywiseScaleBias(
x,
weights=self.bn_weights[2],
name='{0}_bn_scalebias_instance{1}'.format(self.name,
self.instance),
data_layout=self.data_layout)
return lbann.Relu(x,
name='{0}_relu_instance{1}'.format(
self.name, self.instance),
data_layout=self.data_layout)
def make_model(num_vertices=None,
node_features=None,
num_classes=None,
dataset=None,
kernel_type='GCN',
callbacks=None,
num_epochs=1):
'''Construct a model DAG using one of the Graph Kernels
Args:
num_vertices (int): Number of vertices of each graph (default: None)
node_features (int): Number of features per noded (default: None)
num_classes (int): Number of classes as targets (default: None)
dataset (str): Preset data set to use. Either a datset parameter has to be
supplied or all of num_vertices, node_features, and
num_classes have to be supplied. (default: None)
kernel_type (str): Graph Kernel to use in model. Expected one of
GCN, GIN, Graph, or GatedGraph (deafult: GCN)
callbacks (list): Callbacks for the model. If set to None the model description,
GPU usage, training_output, and timer is reported.
(default: None)
num_epochs (int): Number of epochs to run (default: 1)
Returns:
(lbann Model Object: A model object with the supplied callbacks, dataset
presets, and graph kernels.
'''
assert num_vertices != dataset #Ensure atleast one of the values is set
if dataset is not None:
assert num_vertices is None
if dataset == 'MNIST':
num_vertices = 75
num_classes = 10
node_features = 1
elif dataset == 'PROTEINS':
num_vertices = 100
num_classes = 2
node_features = 3
else:
raise Exception("Unkown Dataset")
assert num_vertices is not None
assert num_classes is not None
assert node_features is not None
#----------------------------------
# Reshape and Slice Input Tensor
#----------------------------------
input_ = lbann.Input(target_mode='classification')
# Input dimensions should be (num_vertices * node_features + num_vertices^2 + num_classes )
# Input should have atleast two children since the target is classification
data = lbann_Graph_Data(input_, num_vertices, node_features, num_classes)
feature_matrix = data.x
adj_matrix = data.adj
target = data.y
#----------------------------------
# Perform Graph Convolution
#----------------------------------
if kernel_type == 'GIN':
x = GINConvLayer(feature_matrix, adj_matrix)
elif kernel_type == 'GCN':
x = GCNConvLayer(feature_matrix, adj_matrix)
elif kernel_type == 'Graph':
x = GraphConvLayer(feature_matrix, adj_matrix)
elif kernel_type == 'GatedGraph':
x = GATConvLayer(feature_matrix, adj_matrix)
else:
ValueError(
'Invalid Graph kernel specifier "{}" recieved. Expected one of:\
GIN,GCN,Graph or GatedGraph'.format(kernel_type))
out_channel = x.shape[1]
#----------------------------------
# Apply Reduction on Node Features
#----------------------------------
average_vector = lbann.Constant(value=1 / num_vertices,
num_neurons=str_list([1, num_vertices]),
name="Average_Vector")
x = x.get_mat(out_channel)
x = lbann.MatMul(average_vector, x, name="Node_Feature_Reduction")
# X is now a vector with output_channel dimensions
x = lbann.Reshape(x, dims=str_list([out_channel]), name="Squeeze")
x = lbann.FullyConnected(x, num_neurons=64, name="hidden_layer_1")
x = lbann.Relu(x, name="hidden_layer_1_activation")
x = lbann.FullyConnected(x,
num_neurons=num_classes,
name="Output_Fully_Connected")
#----------------------------------
# Loss Function and Accuracy s
#----------------------------------
probs = lbann.Softmax(x, name="Softmax")
loss = lbann.CrossEntropy(probs, target, name="Cross_Entropy_Loss")
accuracy = lbann.CategoricalAccuracy(probs, target, name="Accuracy")
layers = lbann.traverse_layer_graph(input_)
if callbacks is None:
print_model = lbann.CallbackPrintModelDescription(
) #Prints initial Model after Setup
training_output = lbann.CallbackPrint(
interval=1,
print_global_stat_only=False) #Prints training progress
gpu_usage = lbann.CallbackGPUMemoryUsage()
timer = lbann.CallbackTimer()
callbacks = [print_model, training_output, gpu_usage, timer]
else:
if isinstance(callbacks, list):
callbacks = callbacks
metrics = [lbann.Metric(accuracy, name='accuracy', unit="%")]
model = lbann.Model(num_epochs,
layers=layers,
objective_function=loss,
metrics=metrics,
callbacks=callbacks)
return model
def make_model(NUM_NODES,
NUM_EDGES,
NUM_NODES_FEATURES,
NUM_EDGE_FEATURES,
EMBEDDING_DIM,
EDGE_EMBEDDING_DIM,
NUM_OUT_FEATURES,
NUM_EPOCHS,
NUM_GROUPS=0):
""" Creates an LBANN model for the OGB-LSC PPQM4M Dataset
Args:
NUM_NODES (int): The number of nodes in the largest graph in the dataset (51 for LSC-PPQM4M)
NUM_EDGES (int): The number of edges in the largest graph in the dataset (118 for LSC-PPQM4M)
NUM_NODES_FEATURES (int): The dimensionality of the input node features vector (9 for LSC-PPQM4M)
NUM_EDGE_FEATURES (int): The dimensionality of the input edge feature vectors (3 for LSC-PPQM4M)
EMBEDDING_DIM (int): The embedding dimensionality of the node feature vector
EDGE_EMBEDDING_DIM (int): The embedding dimensionality of the edge feature vector
NUM_OUT_FEATURES (int): The dimensionality of the node feature vectors after graph convolutions
NUM_EPOCHS (int): The number of epochs to train the network
NUM_GROUPS (int): The number of channel groups for distconv channelwise fully connected layer (default : 0)
Returns:
(Model): lbann model object
"""
in_channel = EMBEDDING_DIM
out_channel = NUM_OUT_FEATURES
output_dimension = 1
_input = lbann.Input(data_field='samples')
node_feature_mat, neighbor_feature_mat, edge_feature_mat, edge_indices, target = \
graph_data_splitter(_input,
NUM_NODES,
NUM_EDGES,
NUM_NODES_FEATURES,
NUM_EDGE_FEATURES,
EMBEDDING_DIM,
EDGE_EMBEDDING_DIM)
x = NNConvLayer(node_feature_mat, neighbor_feature_mat, edge_feature_mat,
edge_indices, in_channel, out_channel, EDGE_EMBEDDING_DIM,
NUM_NODES, NUM_EDGES, NUM_GROUPS)
for i, num_neurons in enumerate([256, 128, 32, 8], 1):
x = lbann.FullyConnected(x,
num_neurons=num_neurons,
name="hidden_layer_{}".format(i))
x = lbann.Relu(x, name='hidden_layer_{}_activation'.format(i))
x = lbann.FullyConnected(x, num_neurons=output_dimension, name="output")
loss = lbann.MeanAbsoluteError(x, target)
layers = lbann.traverse_layer_graph(_input)
training_output = lbann.CallbackPrint(interval=1,
print_global_stat_only=False)
gpu_usage = lbann.CallbackGPUMemoryUsage()
timer = lbann.CallbackTimer()
callbacks = [training_output, gpu_usage, timer]
model = lbann.Model(NUM_EPOCHS,
layers=layers,
objective_function=loss,
callbacks=callbacks)
return model
# Input data
input_ = lbann.Input(target_mode='classification')
images = lbann.Identity(input_)
labels = lbann.Identity(input_)
# LeNet
x = lbann.Convolution(images,
num_dims=2,
num_output_channels=6,
num_groups=1,
conv_dims_i=5,
conv_strides_i=1,
conv_dilations_i=1,
has_bias=True)
x = lbann.Relu(x)
x = lbann.Pooling(x,
num_dims=2,
pool_dims_i=2,
pool_strides_i=2,
pool_mode="max")
x = lbann.Convolution(x,
num_dims=2,
num_output_channels=16,
num_groups=1,
conv_dims_i=5,
conv_strides_i=1,
conv_dilations_i=1,
has_bias=True)
x = lbann.Relu(x)
x = lbann.Pooling(x,
args = parser.parse_args()
# Start of layers
# Construct layer graph
input_ = lbann.Input(name='data')
image = lbann.Split(input_, name='image')
dummy = lbann.Dummy(input_, name='dummy')
# Encoder
encode1 = lbann.FullyConnected(image,
name="encode1",
num_neurons=1000,
has_bias=True)
encode1neuron = lbann.Relu(encode1, name="encode1neuron")
encode2 = lbann.FullyConnected(encode1neuron,
name="encode2",
num_neurons=500,
has_bias=True)
encode2neuron = lbann.Relu(encode2, name="encode2neuron")
encode3 = lbann.FullyConnected(encode2neuron,
name="encode3",
num_neurons=250,
has_bias=True)
encode3neuron = lbann.Relu(encode3, name="encode3neuron")
print("Dataset must be cifar10 or imagenet. Try again.")
exit()
# Construct layer graph
input_ = lbann.Input(name='input')
image = lbann.Identity(input_, name='images')
dummy = lbann.Dummy(input_, name='labels')
# Encoder
encode1 = lbann.FullyConnected(image,
name="encode1",
data_layout="model_parallel",
num_neurons=1000,
has_bias=True)
relu1 = lbann.Relu(encode1, name="relu1", data_layout="model_parallel")
dropout1 = lbann.Dropout(relu1,
name="dropout1",
data_layout="model_parallel",
keep_prob=0.8)
decode1 = lbann.FullyConnected(dropout1,
name="decode1",
data_layout="model_parallel",
hint_layer=image,
has_bias=True)
reconstruction = lbann.Sigmoid(decode1,
name="reconstruction",
data_layout="model_parallel")
# Start of layers
# Construct layer graph
input_ = lbann.Input(name='data')
finetunedata = lbann.Split(input_, name='finetunedata')
label = lbann.Split(input_, name='label')
# Encoder
encode1 = lbann.FullyConnected(finetunedata,
name="encode1",
data_layout="model_parallel",
num_neurons=2000,
has_bias=True)
relu1 = lbann.Relu(encode1, name="relu1", data_layout="model_parallel")
encode2 = lbann.FullyConnected(relu1,
name="encode2",
data_layout="model_parallel",
num_neurons=1000,
has_bias=True)
relu2 = lbann.Relu(encode2, name="relu2", data_layout="model_parallel")
encode3 = lbann.FullyConnected(relu2,
name="encode3",
data_layout="model_parallel",
num_neurons=500,
has_bias=True)
def densenet(statistics_group_size, version, cumulative_layer_num,
images_node):
if version == 121:
growth_rate = 32 # k in the paper
layers_per_block = (6, 12, 24, 16)
num_initial_features = 64
elif version == 161:
growth_rate = 48 # k in the paper
layers_per_block = (96, 48, 36, 24)
num_initial_features = 96
else:
raise Exception('Invalid version={v}.'.format(v=version))
batch_norm_size = 4
parent_node, cumulative_layer_num = initial_layer(statistics_group_size,
cumulative_layer_num,
images_node,
num_initial_features)
num_features = num_initial_features
# Start counting dense blocks at 1.
for current_block_num, num_layers in enumerate(layers_per_block, 1):
parent_nodes, cumulative_layer_num = dense_block(
statistics_group_size,
cumulative_layer_num,
parent_node,
batch_norm_size=batch_norm_size,
current_block_num=current_block_num,
growth_rate=growth_rate,
num_layers=num_layers,
num_initial_channels=num_initial_features)
# num_features += num_layers * growth_rate
for node in parent_nodes[1:]:
num_features += node.num_output_channels
parent_node = lbann.Concatenation(parent_nodes)
cumulative_layer_num += 1
log('densenet Concatenation. cumulative_layer_num={n}'.format(
b=current_block_num, n=cumulative_layer_num))
if current_block_num != len(layers_per_block):
parent_node, cumulative_layer_num = transition_layer(
statistics_group_size,
current_block_num,
cumulative_layer_num,
parent_node,
# In Python 3, this is integer division.
num_output_channels=num_features // 2,
)
num_features //= 2
batch_normalization_node = standard_batchnorm(statistics_group_size,
parent_node)
cumulative_layer_num += 1
log('densenet BatchNormalization. cumulative_layer_num={n}'.format(
b=current_block_num, n=cumulative_layer_num))
relu_node = lbann.Relu(batch_normalization_node)
cumulative_layer_num += 1
log('densenet Relu. cumulative_layer_num={n}'.format(
b=current_block_num, n=cumulative_layer_num))
probs = classification_layer(cumulative_layer_num, relu_node)
return probs
def forward(self, x, memory, src_mask=None, tgt_mask=None):
"""Apply Transformer decoder layer.
Args:
x (lbann.Layer): Sequence of input vectors.
memory (lbann.Layer): Sequence of vectors produced by
Transformer encoder stack.
src_mask (lbann.Layer, optional): Attention mask for
second attention module (attends to both `x` and
`memory`).
tgt_mask (lbann.Layer, optional): Attention mask for first
attention module (attends only to `x`).
Returns:
lbann.Layer: Sequence of output vectors.
"""
self.instance += 1
name = f'{self.name}_instance{self.instance}'
# Self-attention with residual connection
y = self.attention1(x, x, x, mask=tgt_mask)
if self.dropout_prob > 0:
y = lbann.Dropout(
y,
keep_prob=1 - self.dropout_prob,
name=f'{name}_drop1',
)
z = lbann.Sum(x, y, name=f'{name}_sum1')
z = lbann.InstanceNorm(z, name=f'{name}_norm1')
x = z
# Attention on encoder output with residual connection
y = self.attention2(x, memory, memory, mask=src_mask)
if self.dropout_prob > 0:
y = lbann.Dropout(
y,
keep_prob=1 - self.dropout_prob,
name=f'{name}_drop2',
)
z = lbann.Sum(x, y, name=f'{name}_sum2')
z = lbann.InstanceNorm(z, name=f'{name}_norm2')
x = z
# Feedforward network with residual connection
y = lbann.ChannelwiseFullyConnected(
x,
weights=self.fc1_weights,
output_channel_dims=[self.feedforward_dim],
name=f'{name}_fc1',
)
y = lbann.Relu(y, name=f'{name}_relu1')
if self.dropout_prob > 0:
y = lbann.Dropout(
y,
keep_prob=1 - self.dropout_prob,
name=f'{name}_drop3',
)
y = lbann.ChannelwiseFullyConnected(
y,
weights=self.fc2_weights,
output_channel_dims=[self.embed_dim],
name=f'{name}_fc2',
)
if self.dropout_prob > 0:
y = lbann.Dropout(
y,
keep_prob=1 - self.dropout_prob,
name=f'{name}_drop4',
)
z = lbann.Sum(x, y, name=f'{name}_sum3')
z = lbann.InstanceNorm(z, name=f'{name}_norm3')
return z
# Start of layers
# Construct layer graph
input_ = lbann.Input(name='input', target_mode="reconstruction")
data = lbann.Split(input_, name='data')
dummy = lbann.Dummy(input_, name='dummy')
# Encoder
encode1 = lbann.FullyConnected(data,
name="encode1",
data_layout="model_parallel",
num_neurons=2000,
has_bias=True)
relu1 = lbann.Relu(encode1, name="relu1", data_layout="model_parallel")
encode2 = lbann.FullyConnected(relu1,
name="encode2",
data_layout="model_parallel",
num_neurons=1000,
has_bias=True)
relu2 = lbann.Relu(encode2, name="relu2", data_layout="model_parallel")
encode3 = lbann.FullyConnected(relu2,
name="encode3",
data_layout="model_parallel",
num_neurons=500,
has_bias=True)
def make_model(num_vertices=None,
node_features=None,
num_classes=None,
kernel_type='GCN',
callbacks=None,
num_epochs=1):
'''Construct a model DAG using one of the Graph Kernels
Args:
num_vertices (int): Number of vertices of each graph (default: None)
node_features (int): Number of features per noded (default: None)
num_classes (int): Number of classes as targets (default: None)
kernel_type (str): Graph Kernel to use in model. Expected one of
GCN, GIN, Graph, or GatedGraph (deafult: GCN)
callbacks (list): Callbacks for the model. If set to None the model description,
GPU usage, training_output, and timer is reported.
(default: None)
num_epochs (int): Number of epochs to run (default: 1)
Returns:
(lbann.Model) : A model object with the supplied callbacks, dataset
presets, and graph kernels.
'''
num_vertices = 100
num_classes = 2
node_feature_size = 3
max_edges = 415
#----------------------------------
# Reshape and Slice Input Tensor
#----------------------------------
input_ = lbann.Input(data_field='samples')
# Input dimensions should be (num_vertices * node_features + num_vertices^2 + num_classes )
data = Graph_Data_Parser(input_, num_vertices, node_feature_size,
max_edges, num_classes)
feature_matrix = data['node_features']
source_indices = data['source_indices']
target_indices = data['target_indices']
target = data['target']
#----------------------------------
# Select Graph Convolution
#----------------------------------
output_channels = 16
graph_kernel_op = None
if kernel_type == 'GIN':
graph_kernel_op = GINConvLayer
elif kernel_type == 'GCN':
graph_kernel_op = GCNConvLayer
elif kernel_type == 'Graph':
graph_kernel_op = GraphConvLayer
elif kernel_type == 'GatedGraph':
graph_kernel_op = GATConvLayer
else:
raise ValueError(
'Invalid Graph kernel specifier "{}" recieved. Expected one of:\
GIN,GCN,Graph or GatedGraph'.format(kernel_type))
#----------------------------------
# Perform Graph Convolution
#----------------------------------
x = graph_kernel_op(feature_matrix, source_indices, target_indices,
num_vertices, max_edges, node_feature_size,
output_channels)
#----------------------------------
# Apply Reduction on Node Features
#----------------------------------
average_vector = lbann.Constant(value=1 / num_vertices,
num_neurons=str_list([1, num_vertices]),
name="Average_Vector")
x = lbann.MatMul(average_vector, x, name="Node_Feature_Reduction")
# X is now a vector with output_channel dimensions
x = lbann.Reshape(x, dims=str_list([output_channels]), name="Squeeze")
x = lbann.FullyConnected(x, num_neurons=64, name="hidden_layer_1")
x = lbann.Relu(x, name="hidden_layer_1_activation")
x = lbann.FullyConnected(x,
num_neurons=num_classes,
name="Output_Fully_Connected")
#----------------------------------
# Loss Function and Accuracy s
#----------------------------------
probs = lbann.Softmax(x, name="Softmax")
loss = lbann.CrossEntropy(probs, target, name="Cross_Entropy_Loss")
accuracy = lbann.CategoricalAccuracy(probs, target, name="Accuracy")
layers = lbann.traverse_layer_graph(input_)
if callbacks is None:
print_model = lbann.CallbackPrintModelDescription(
) #Prints initial Model after Setup
training_output = lbann.CallbackPrint(
interval=1,
print_global_stat_only=False) #Prints training progress
gpu_usage = lbann.CallbackGPUMemoryUsage()
timer = lbann.CallbackTimer()
callbacks = [print_model, training_output, gpu_usage, timer]
else:
if isinstance(callbacks, list):
callbacks = callbacks
metrics = [lbann.Metric(accuracy, name='accuracy', unit="%")]
model = lbann.Model(num_epochs,
layers=layers,
objective_function=loss,
metrics=metrics,
callbacks=callbacks)
return model
