def evaluate_model(_):
if FLAGS.gpu >= 0:
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = str(FLAGS.gpu)
else:
logging.info(
"Not using a GPU. Specify one by passing --gpu {idx} with idx >= 0."
)
if FLAGS.dataset == "hdataset":
data_index = hdataset_provider.HDatasetGenerator.initialize_hdataset(
FLAGS.data)
eval_index = data_index[~data_index['train']]
eval_generator = hdataset_provider.HDatasetGenerator(eval_index)
else:
data_index = provider.ShapeNetGenerator.initialize_shapenet(
FLAGS.data, class_choice=FLAGS.categories)
eval_index = data_index[~data_index['train']]
eval_generator = provider.ShapeNetGenerator(eval_index,
n_points=FLAGS.n_points)
eval_dataset = eval_generator.create_dataset().batch(FLAGS.batch_size)
atlas_net = model.AtlasNetV2(structure_type=FLAGS.structure_type,
adjustment_type=FLAGS.adjustment)
logging.info(f"Using {len(eval_generator)} samples for evaluation.")
atlas_net.compile(optimizer=tf.keras.optimizers.Adam(),
loss=loss.chamfer_loss,
metrics=['mean_squared_error', loss.mean_chamfer_loss])
atlas_net.predict(eval_dataset, steps=1)
atlas_net.load_weights(FLAGS.model)
results = atlas_net.evaluate(eval_dataset,
steps=len(eval_generator) // FLAGS.batch_size,
return_dict=True)
results['model_metadata'] = {
'num_params': atlas_net.count_params(),
'structure_type': FLAGS.structure_type,
'adjustment_type': FLAGS.adjustment,
}
results_json = json.dumps(results)
with tf.io.gfile.GFile(FLAGS.result_file, "w") as outfile:
json.dump(results_json, outfile)
def test_load_saved_model(self):
atlasnet = model.AtlasNetV2(structure_type='patch',
adjustment_type='linear',
n_structures=10)
atlasnet.compile(optimizer='Adam',
loss=loss.chamfer_loss,
metrics=['mean_squared_error'])
df = provider.ShapeNetGenerator.initialize_shapenet(
"../data/shape_net/", class_choice=['plane'], train_frac=0.7)
generator = provider.ShapeNetGenerator(df,
n_points=10000,
visualize=False)
dataset = generator.create_dataset().batch(BATCH_SIZE)
atlasnet.evaluate(dataset, steps=1)
atlasnet.load_weights('weights/atlasnet')
atlasnet.evaluate(dataset, steps=1)
def test_fit_with_real_input_and_serialization(self):
df = provider.ShapeNetGenerator.initialize_shapenet(
"../data/shape_net/", class_choice=['plane'], train_frac=0.7)
atlasnet = model.AtlasNetV2(structure_type='point',
adjustment_type='linear',
n_structures=10)
generator = provider.ShapeNetGenerator(df,
n_points=10000,
visualize=False)
dataset = generator.create_dataset().batch(BATCH_SIZE)
atlasnet.compile(optimizer='Adam',
loss=loss.chamfer_loss,
metrics=['mean_squared_error'])
print("Model compiled successfully!")
atlasnet.fit(dataset, epochs=1, steps_per_epoch=1)
atlasnet.save_weights('weights/atlasnet')
import json
print(json.dumps(atlasnet.get_config()))
def test_shapenet_provider_batch_size(self):
dataset = provider.ShapeNetGenerator.initialize_shapenet("../data/shape_net/", class_choice=None)
generator = provider.ShapeNetGenerator(dataset, n_points=42)
data = iter(generator.create_dataset().batch(8))
assert next(data)[0].shape == (8, 42, 3)
def test_shapenet_provider_len(self):
dataset = provider.ShapeNetGenerator.initialize_shapenet("../data/shape_net/", class_choice=['plane'])
generator = provider.ShapeNetGenerator(dataset, n_points=42)
assert len(dataset) == len(generator)
def train_model(_):
if FLAGS.fromcheckpoint is None:
now = datetime.now()
mlflow.set_experiment(
f"{FLAGS.experiment} vom {now.strftime('%m/%d/%Y, %H:%M:%S')}")
else:
mlflow.set_experiment(FLAGS.experiment)
with mlflow.start_run():
if FLAGS.gpu >= 0:
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = str(FLAGS.gpu)
else:
logging.info(
"Not using a GPU. Specify one by passing --gpu {idx} with idx >= 0."
)
data_index = provider.ShapeNetGenerator.initialize_shapenet(
FLAGS.data, class_choice=FLAGS.categories)
train_index = data_index[data_index['train']]
train_files, val_files = random_split(train_index, 1 - FLAGS.val_split)
train_generator = provider.ShapeNetGenerator(train_files,
n_points=FLAGS.n_points)
eval_generator = provider.ShapeNetGenerator(val_files,
n_points=FLAGS.n_points)
train_dataset = train_generator.create_dataset().batch(
FLAGS.batch_size)
eval_dataset = eval_generator.create_dataset().batch(FLAGS.batch_size)
logging.info(
f"Using {len(train_generator)} samples for training and {len(eval_generator)} samples for evaluation."
)
atlas_net = model.AtlasNetV2(structure_type=FLAGS.structure_type,
adjustment_type=FLAGS.adjustment)
callbacks = [
tf.keras.callbacks.TensorBoard(FLAGS.logdir, write_graph=False),
tf.keras.callbacks.ModelCheckpoint(FLAGS.checkpoint,
save_weights_only=True,
save_best_only=True),
MLflowCallback(),
]
atlas_net.compile(
optimizer=tf.keras.optimizers.Adam(FLAGS.learning_rate),
loss=loss.chamfer_loss,
metrics=['mean_squared_error', loss.mean_chamfer_loss])
if FLAGS.fromcheckpoint is not None:
atlas_net.fit(train_dataset, steps_per_epoch=1, epochs=1)
atlas_net.load_weights(FLAGS.from_checkpoint)
atlas_net.fit(train_dataset,
steps_per_epoch=len(train_generator) // FLAGS.batch_size,
epochs=FLAGS.epochs,
callbacks=callbacks,
validation_data=eval_dataset,
validation_steps=len(eval_generator) // FLAGS.batch_size)
mlflow.log_param('AtlasNet-Config', atlas_net.get_config())
atlas_net.save_weights(FLAGS.output)