def set_up_experiment(args,
input_,
probs,
labels):
# Set up objective function
cross_entropy = lbann.CrossEntropy([probs, labels])
layers = list(lbann.traverse_layer_graph(input_))
weights = set()
for l in layers:
weights.update(l.weights)
# scale = weight decay
l2_reg = lbann.L2WeightRegularization(weights=weights, scale=1e-4)
objective_function = lbann.ObjectiveFunction([cross_entropy, l2_reg])
# Set up model
top1 = lbann.CategoricalAccuracy([probs, labels])
top5 = lbann.TopKCategoricalAccuracy([probs, labels], k=5)
metrics = [lbann.Metric(top1, name='top-1 accuracy', unit='%'),
lbann.Metric(top5, name='top-5 accuracy', unit='%')]
callbacks = [lbann.CallbackPrint(),
lbann.CallbackTimer(),
lbann.CallbackDropFixedLearningRate(
drop_epoch=[30, 60], amt=0.1)]
model = lbann.Model(args.mini_batch_size,
args.num_epochs,
layers=layers,
weights=weights,
objective_function=objective_function,
metrics=metrics,
callbacks=callbacks)
# Load data reader from prototext
data_reader_proto = lbann.lbann_pb2.LbannPB()
with open(args.data_reader, 'r') as f:
txtf.Merge(f.read(), data_reader_proto)
data_reader_proto = data_reader_proto.data_reader
# Set up optimizer
if args.optimizer == 'sgd':
print('Creating sgd optimizer')
optimizer = lbann.optimizer.SGD(
learn_rate=args.optimizer_learning_rate,
momentum=0.9,
nesterov=True
)
else:
optimizer = lbann.contrib.args.create_optimizer(args)
# Save prototext to args.prototext
if args.prototext:
lbann.proto.save_prototext(args.prototext,
model=model,
optimizer=optimizer,
data_reader=data_reader_proto)
return model, data_reader_proto, optimizer
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 set_up_experiment(args, input_, probs, labels):
# Set up objective function
cross_entropy = lbann.CrossEntropy([probs, labels])
layers = list(lbann.traverse_layer_graph(input_))
l2_reg_weights = set()
for l in layers:
if type(l) == lbann.Convolution or type(l) == lbann.FullyConnected:
l2_reg_weights.update(l.weights)
# scale = weight decay
l2_reg = lbann.L2WeightRegularization(weights=l2_reg_weights, scale=1e-4)
objective_function = lbann.ObjectiveFunction([cross_entropy, l2_reg])
# Set up model
top1 = lbann.CategoricalAccuracy([probs, labels])
top5 = lbann.TopKCategoricalAccuracy([probs, labels], k=5)
metrics = [
lbann.Metric(top1, name='top-1 accuracy', unit='%'),
lbann.Metric(top5, name='top-5 accuracy', unit='%')
]
callbacks = [
lbann.CallbackPrint(),
lbann.CallbackTimer(),
lbann.CallbackDropFixedLearningRate(drop_epoch=[30, 60], amt=0.1)
]
model = lbann.Model(args.num_epochs,
layers=layers,
objective_function=objective_function,
metrics=metrics,
callbacks=callbacks)
# Set up data reader
data_reader = data.imagenet.make_data_reader(num_classes=args.num_classes)
# Set up optimizer
if args.optimizer == 'sgd':
print('Creating sgd optimizer')
optimizer = lbann.optimizer.SGD(
learn_rate=args.optimizer_learning_rate,
momentum=0.9,
nesterov=True)
else:
optimizer = lbann.contrib.args.create_optimizer(args)
# Setup trainer
trainer = lbann.Trainer(mini_batch_size=args.mini_batch_size)
return trainer, model, data_reader, optimizer
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 setup(data_reader_file,
name='classifier',
num_labels=200,
mini_batch_size=128,
num_epochs=1000,
learning_rate=0.1,
bn_statistics_group_size=2,
fc_data_layout='model_parallel',
warmup_epochs=50,
learning_rate_drop_interval=50,
learning_rate_drop_factor=0.25,
checkpoint_interval=None):
# Setup input data
input = lbann.Input(target_mode = 'classification')
images = lbann.Identity(input)
labels = lbann.Identity(input)
# Classification network
head_cnn = modules.ResNet(bn_statistics_group_size=bn_statistics_group_size)
class_fc = lbann.modules.FullyConnectedModule(num_labels,
activation=lbann.Softmax,
name=f'{name}_fc',
data_layout=fc_data_layout)
x = head_cnn(images)
probs = class_fc(x)
# Setup objective function
cross_entropy = lbann.CrossEntropy([probs, labels])
l2_reg_weights = set()
for l in lbann.traverse_layer_graph(input):
if type(l) == lbann.Convolution or type(l) == lbann.FullyConnected:
l2_reg_weights.update(l.weights)
l2_reg = lbann.L2WeightRegularization(weights=l2_reg_weights, scale=0.0002)
obj = lbann.ObjectiveFunction([cross_entropy, l2_reg])
# Setup model
metrics = [lbann.Metric(lbann.CategoricalAccuracy([probs, labels]),
name='accuracy', unit='%')]
callbacks = [lbann.CallbackPrint(), lbann.CallbackTimer()]
if checkpoint_interval:
callbacks.append(
lbann.CallbackCheckpoint(
checkpoint_dir='ckpt',
checkpoint_epochs=5
)
)
# Learning rate schedules
if warmup_epochs:
callbacks.append(
lbann.CallbackLinearGrowthLearningRate(
target=learning_rate * mini_batch_size / 128,
num_epochs=warmup_epochs
)
)
if learning_rate_drop_factor:
callbacks.append(
lbann.CallbackDropFixedLearningRate(
drop_epoch=list(range(0, num_epochs, learning_rate_drop_interval)),
amt=learning_rate_drop_factor)
)
# Construct model
model = lbann.Model(num_epochs,
layers=lbann.traverse_layer_graph(input),
objective_function=obj,
metrics=metrics,
callbacks=callbacks)
# Setup optimizer
# opt = lbann.Adam(learn_rate=learning_rate, beta1=0.9, beta2=0.999, eps=1e-8)
opt = lbann.SGD(learn_rate=learning_rate, momentum=0.9)
# Load data reader from prototext
data_reader_proto = lbann.lbann_pb2.LbannPB()
with open(data_reader_file, 'r') as f:
google.protobuf.text_format.Merge(f.read(), data_reader_proto)
data_reader_proto = data_reader_proto.data_reader
for reader_proto in data_reader_proto.reader:
reader_proto.python.module_dir = os.path.dirname(os.path.realpath(__file__))
# Return experiment objects
return model, data_reader_proto, opt
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.FullyConnected(x, num_neurons=120, has_bias=True)
x = lbann.Relu(x)
x = lbann.FullyConnected(x, num_neurons=84, has_bias=True)
x = lbann.Relu(x)
x = lbann.FullyConnected(x, num_neurons=10, has_bias=True)
probs = lbann.Softmax(x)
# Loss function and accuracy
loss = lbann.CrossEntropy(probs, labels)
acc = lbann.CategoricalAccuracy(probs, labels)
# ----------------------------------
# Setup experiment
# ----------------------------------
# Setup model
mini_batch_size = 64
num_epochs = 20
model = lbann.Model(num_epochs,
layers=lbann.traverse_layer_graph(input_),
objective_function=loss,
metrics=[lbann.Metric(acc, name='accuracy', unit='%')],
callbacks=[
lbann.CallbackPrintModelDescription(),
lbann.contrib.args.add_optimizer_arguments(
parser,
default_optimizer="adam",
default_learning_rate=0.001,
)
args = parser.parse_args()
parallel_strategy = get_parallel_strategy_args(
sample_groups=args.mini_batch_size, depth_groups=args.depth_groups)
# Construct layer graph
input = lbann.Input(target_mode='label_reconstruction')
volume = lbann.Identity(input)
output = UNet3D()(volume)
segmentation = lbann.Identity(input)
ce = lbann.CrossEntropy([output, segmentation], use_labels=True)
obj = lbann.ObjectiveFunction([ce])
layers = list(lbann.traverse_layer_graph(input))
for l in layers:
l.parallel_strategy = parallel_strategy
# Setup model
metrics = [lbann.Metric(ce, name='CE', unit='')]
callbacks = [
lbann.CallbackPrint(),
lbann.CallbackTimer(),
lbann.CallbackGPUMemoryUsage(),
lbann.CallbackProfiler(skip_init=True),
]
# # TODO: Use polynomial learning rate decay (https://github.com/LLNL/lbann/issues/1581)
# callbacks.append(
pred_fc = lbann.modules.FullyConnectedModule(vocab_size,
data_layout='model_parallel')
# Iterate through RNN steps
loss = []
for step in range(sequence_length - 1):
# Predict next token with RNN
x = embeddings_list[step]
x, lstm_state = lstm(x, lstm_state)
x = pred_fc(x)
pred = lbann.Softmax(x)
# Evaluate prediction with cross entropy
ground_truth = lbann.OneHot(tokens_list[step + 1], size=vocab_size)
cross_entropy = lbann.CrossEntropy([pred, ground_truth])
loss.append(lbann.LayerTerm(cross_entropy,
scale=1 / (sequence_length - 1)))
# ----------------------------------
# Create data reader
# ----------------------------------
reader = lbann.reader_pb2.DataReader()
_reader = reader.reader.add()
_reader.name = 'python'
_reader.role = 'train'
_reader.shuffle = True
_reader.percent_of_data_to_use = 1.0
_reader.python.module = 'dataset'
_reader.python.module_dir = current_dir
imagenet_labels,
bn_statistics_group_size=args.bn_statistics_group_size)
else:
# Some other Wide ResNet.
resnet = resnet_variant_dict[args.resnet](
imagenet_labels,
bn_statistics_group_size=args.bn_statistics_group_size,
width=args.width)
# Construct layer graph
input_ = lbann.Input(target_mode='classification')
images = lbann.Identity(input_)
labels = lbann.Identity(input_)
preds = resnet(images)
probs = lbann.Softmax(preds)
cross_entropy = lbann.CrossEntropy(probs, labels)
top1 = lbann.CategoricalAccuracy(probs, labels)
top5 = lbann.TopKCategoricalAccuracy(probs, labels, k=5)
layers = list(lbann.traverse_layer_graph(input_))
# Setup tensor core operations (just to demonstrate enum usage)
tensor_ops_mode = lbann.ConvTensorOpsMode.NO_TENSOR_OPS
for l in layers:
if type(l) == lbann.Convolution:
l.conv_tensor_op_mode = tensor_ops_mode
# Setup objective function
l2_reg_weights = set()
for l in layers:
if type(l) == lbann.Convolution or type(l) == lbann.FullyConnected:
l2_reg_weights.update(l.weights)
data_layout="model_parallel",
num_neurons=100,
has_bias=True)
relu5 = lbann.Relu(encode5, name="relu5", data_layout="model_parallel")
ip2 = lbann.FullyConnected(relu5,
name="ip2",
data_layout="model_parallel",
num_neurons=2,
has_bias=True)
prob = lbann.Softmax(ip2, name="prob", data_layout="model_parallel")
cross_entropy = lbann.CrossEntropy([prob, label],
name="cross_entropy",
data_layout="model_parallel")
categorical_accuracy = lbann.CategoricalAccuracy([prob, label],
name="categorical_accuracy",
data_layout="model_parallel")
layer_list = list(lbann.traverse_layer_graph(input_))
# Set up objective function
layer_term = lbann.LayerTerm(cross_entropy)
obj = lbann.ObjectiveFunction(layer_term)
# Metrics
metrics = [lbann.Metric(categorical_accuracy, name="accuracy")]
def make_model(
num_epochs,
embed_dim,
num_heads,
label_smoothing,
):
# Embedding weights
var = 2 / (embed_dim + vocab_size) # Glorot initialization
embedding_weights = lbann.Weights(
name='embeddings',
initializer=lbann.NormalInitializer(standard_deviation=math.sqrt(var)),
)
# Input is two sequences of token IDs
input_ = lbann.Input(data_field='samples')
# Get sequences of embedding vectors
# Note: Scale embeddings by sqrt(embed_dim).
# Note: Decoder input is shifted right, so embedding for last
# token isn't needed.
embeddings_tokens = lbann.Identity(
lbann.Slice(
input_,
axis=0,
slice_points=str_list([0, 2 * sequence_length - 1]),
))
embeddings = lbann.Embedding(
embeddings_tokens,
weights=embedding_weights,
num_embeddings=vocab_size,
embedding_dim=embed_dim,
padding_idx=pad_index,
)
embeddings = lbann.WeightedSum(
embeddings,
scaling_factors=str(math.sqrt(embed_dim)),
)
embeddings_slice = lbann.Slice(
embeddings,
axis=0,
slice_points=str_list([0, sequence_length, 2 * sequence_length - 1]),
)
encoder_input = lbann.Identity(embeddings_slice)
decoder_input = lbann.Identity(embeddings_slice)
# Apply transformer model
transformer = lbann.models.Transformer(
hidden_size=embed_dim,
num_heads=num_heads,
name='transformer',
)
result = transformer(
encoder_input,
sequence_length,
decoder_input,
sequence_length - 1,
)
# Reconstruct decoder input
preds = lbann.ChannelwiseFullyConnected(
result,
weights=embedding_weights,
output_channel_dims=[vocab_size],
bias=False,
transpose=True,
)
preds = lbann.ChannelwiseSoftmax(preds)
preds = lbann.Slice(preds,
axis=0,
slice_points=str_list(range(sequence_length)))
preds = [lbann.Identity(preds) for _ in range(sequence_length - 1)]
# Count number of non-pad tokens
label_tokens = lbann.Identity(
lbann.Slice(
input_,
slice_points=str_list([sequence_length + 1, 2 * sequence_length]),
))
pads = lbann.Constant(value=pad_index,
num_neurons=str(sequence_length - 1))
is_not_pad = lbann.NotEqual(label_tokens, pads)
num_not_pad = lbann.Reduction(is_not_pad, mode='sum')
# Cross entropy loss with label smoothing
label_tokens = lbann.Slice(
label_tokens,
slice_points=str_list(range(sequence_length)),
)
label_tokens = [
lbann.Identity(label_tokens) for _ in range(sequence_length - 1)
]
if label_smoothing > 0:
uniform_label = lbann.Constant(value=1 / vocab_size,
num_neurons=str_list([1, vocab_size]))
loss = []
for i in range(sequence_length - 1):
label = lbann.OneHot(label_tokens[i], size=vocab_size)
label = lbann.Reshape(label, dims=str_list([1, vocab_size]))
if label_smoothing > 0:
label = lbann.WeightedSum(
label,
uniform_label,
scaling_factors=str_list(
[1 - label_smoothing, label_smoothing]),
)
loss.append(lbann.CrossEntropy(preds[i], label))
loss = lbann.Concatenation(loss)
# Average cross entropy over non-pad tokens
loss_scales = lbann.Divide(
is_not_pad,
lbann.Tessellate(num_not_pad, hint_layer=is_not_pad),
)
loss = lbann.Multiply(loss, loss_scales)
loss = lbann.Reduction(loss, mode='sum')
# Construct model
metrics = []
callbacks = [lbann.CallbackPrint(), lbann.CallbackTimer()]
return lbann.Model(
num_epochs,
layers=lbann.traverse_layer_graph(input_),
objective_function=loss,
metrics=metrics,
callbacks=callbacks,
)
def setup(num_patches=3,
mini_batch_size=512,
num_epochs=75,
learning_rate=0.005,
bn_statistics_group_size=2,
fc_data_layout='model_parallel',
warmup=True,
checkpoint_interval=None):
# Data dimensions
patch_dims = patch_generator.patch_dims
num_labels = patch_generator.num_labels(num_patches)
# Extract tensors from data sample
input = lbann.Input()
slice_points = [0]
for _ in range(num_patches):
patch_size = functools.reduce(operator.mul, patch_dims)
slice_points.append(slice_points[-1] + patch_size)
slice_points.append(slice_points[-1] + num_labels)
sample = lbann.Slice(input, slice_points=str_list(slice_points))
patches = [
lbann.Reshape(sample, dims=str_list(patch_dims))
for _ in range(num_patches)
]
labels = lbann.Identity(sample)
# Siamese network
head_cnn = modules.ResNet(
bn_statistics_group_size=bn_statistics_group_size)
heads = [head_cnn(patch) for patch in patches]
heads_concat = lbann.Concatenation(heads)
# Classification network
class_fc1 = modules.FcBnRelu(
4096,
statistics_group_size=bn_statistics_group_size,
name='siamese_class_fc1',
data_layout=fc_data_layout)
class_fc2 = modules.FcBnRelu(
4096,
statistics_group_size=bn_statistics_group_size,
name='siamese_class_fc2',
data_layout=fc_data_layout)
class_fc3 = lbann.modules.FullyConnectedModule(num_labels,
activation=lbann.Softmax,
name='siamese_class_fc3',
data_layout=fc_data_layout)
x = class_fc1(heads_concat)
x = class_fc2(x)
probs = class_fc3(x)
# Setup objective function
cross_entropy = lbann.CrossEntropy([probs, labels])
l2_reg_weights = set()
for l in lbann.traverse_layer_graph(input):
if type(l) == lbann.Convolution or type(l) == lbann.FullyConnected:
l2_reg_weights.update(l.weights)
l2_reg = lbann.L2WeightRegularization(weights=l2_reg_weights, scale=0.0002)
obj = lbann.ObjectiveFunction([cross_entropy, l2_reg])
# Setup model
metrics = [
lbann.Metric(lbann.CategoricalAccuracy([probs, labels]),
name='accuracy',
unit='%')
]
callbacks = [lbann.CallbackPrint(), lbann.CallbackTimer()]
if checkpoint_interval:
callbacks.append(
lbann.CallbackCheckpoint(checkpoint_dir='ckpt',
checkpoint_epochs=5))
# Learning rate schedules
if warmup:
callbacks.append(
lbann.CallbackLinearGrowthLearningRate(target=learning_rate *
mini_batch_size / 128,
num_epochs=5))
callbacks.append(
lbann.CallbackDropFixedLearningRate(drop_epoch=list(range(0, 100, 15)),
amt=0.25))
# Construct model
model = lbann.Model(num_epochs,
layers=lbann.traverse_layer_graph(input),
objective_function=obj,
metrics=metrics,
callbacks=callbacks)
# Setup optimizer
opt = lbann.SGD(learn_rate=learning_rate, momentum=0.9)
# opt = lbann.Adam(learn_rate=learning_rate, beta1=0.9, beta2=0.999, eps=1e-8)
# Setup data reader
data_reader = make_data_reader(num_patches)
# Return experiment objects
return model, data_reader, opt
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
def construct_model(run_args):
"""Construct LBANN model.
Initial model for ATOM molecular SMILES generation
Network architecture and training hyperparameters from
https://github.com/samadejacobs/moses/tree/master/moses/char_rnn
"""
pad_index = run_args.pad_index
assert pad_index is not None
sequence_length = run_args.sequence_length
assert sequence_length is not None
print("sequence length is {}".format(sequence_length))
data_layout = "data_parallel"
# Layer graph
_input = lbann.Input(name="inp_tensor", data_field='samples')
print(sequence_length)
x_slice = lbann.Slice(
_input,
axis=0,
slice_points=str_list(range(sequence_length + 1)),
name="inp_slice",
)
# embedding layer
emb = []
embedding_dim = run_args.embedding_dim
num_embeddings = run_args.num_embeddings
assert embedding_dim is not None
assert num_embeddings is not None
emb_weights = lbann.Weights(
initializer=lbann.NormalInitializer(mean=0, standard_deviation=1),
name="emb_matrix",
)
lstm1 = lbann.modules.GRU(size=run_args.hidden, data_layout=data_layout)
fc = lbann.modules.FullyConnectedModule(size=num_embeddings,
data_layout=data_layout)
last_output = lbann.Constant(
value=0.0,
num_neurons="{}".format(run_args.hidden),
data_layout=data_layout,
name="lstm_init_output",
)
lstm1_prev_state = [last_output]
loss = []
idl = []
for i in range(sequence_length):
idl.append(
lbann.Identity(x_slice, name="slice_idl_" + str(i), device="CPU"))
for i in range(sequence_length - 1):
emb_l = lbann.Embedding(
idl[i],
name="emb_" + str(i),
weights=emb_weights,
embedding_dim=embedding_dim,
num_embeddings=num_embeddings,
)
x, lstm1_prev_state = lstm1(emb_l, lstm1_prev_state)
fc_l = fc(x)
y_soft = lbann.Softmax(fc_l, name="soft_" + str(i))
gt = lbann.OneHot(idl[i + 1], size=num_embeddings)
ce = lbann.CrossEntropy([y_soft, gt], name="loss_" + str(i))
# mask padding in input
pad_mask = lbann.NotEqual(
[idl[i], lbann.Constant(value=pad_index, num_neurons="1")], )
ce_mask = lbann.Multiply([pad_mask, ce], name="loss_mask_" + str(i))
loss.append(lbann.LayerTerm(ce_mask, scale=1 / (sequence_length - 1)))
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(),
lbann.CallbackStepLearningRate(step=run_args.step_size,
amt=run_args.gamma),
lbann.CallbackDumpWeights(directory=run_args.dump_weights_dir,
epoch_interval=1),
]
# Construct model
return lbann.Model(run_args.num_epochs,
layers=layers,
weights=weights,
objective_function=obj,
callbacks=callbacks)
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.FullyConnected(x, num_neurons=120, has_bias=True)
x = lbann.Relu(x)
x = lbann.FullyConnected(x, num_neurons=84, has_bias=True)
x = lbann.Relu(x)
x = lbann.FullyConnected(x, num_neurons=10, has_bias=True)
probs = lbann.Softmax(x)
# Loss function and accuracy
loss = lbann.CrossEntropy([probs, labels])
acc = lbann.CategoricalAccuracy([probs, labels])
# ----------------------------------
# Setup experiment
# ----------------------------------
# Setup model
mini_batch_size = 64
num_epochs = 20
model = lbann.Model(mini_batch_size,
num_epochs,
layers=lbann.traverse_layer_graph(input),
objective_function=loss,
metrics=[lbann.Metric(acc, name='accuracy', unit='%')],
callbacks=[lbann.CallbackPrint(),
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, or Graph (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.
'''
num_vertices = 100
num_classes = 2
node_features = 3
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(data_field='samples')
# Input dimensions should be (num_vertices * node_features + num_vertices^2 + num_classes )
# input should have atleast two children since the target is classification
sample_dims = num_vertices * node_features + (num_vertices**
2) + num_classes
graph_dims = num_vertices * node_features + (num_vertices**2)
feature_matrix_size = num_vertices * node_features
graph_input = lbann.Slice(input_,
axis=0,
slice_points=str_list([
0, feature_matrix_size, graph_dims,
sample_dims
]),
name="Graph_Input")
feature_matrix = lbann.Reshape(graph_input,
dims=str_list([num_vertices,
node_features]),
name="Node_features")
adj_matrix = lbann.Reshape(graph_input,
dims=str_list([num_vertices, num_vertices]),
name="Adj_Mat")
target = lbann.Identity(graph_input, name="Target")
target = lbann.Reshape(target, dims=str(num_classes))
#----------------------------------
# Perform Graph Convolution
#----------------------------------
if kernel_type == 'GCN':
x = DGCN_layer(feature_matrix, adj_matrix, node_features)
elif kernel_type == 'Graph':
x = DGraph_Layer(feature_matrix, adj_matrix, node_features)
else:
ValueError(
'Invalid Graph kernel specifier "{}" recieved. Expected one of:\
GCN or Graph'.format(kernel_type))
out_channel = 256
#----------------------------------
# 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([out_channel]), name="Squeeze")
x = lbann.FullyConnected(x, num_neurons=256, 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
