示例#1
0
    def test_gnnmodel_build(self):
        nodes = tf.one_hot([2, 4, 1, 3, 3], 16)
        # neighbors -> 5 nodes, 2 neighbors
        nlist = np.zeros((5, 2), dtype=np.int)
        # make each atom neighbors with subsequent 2 atoms, mod 5
        for i in range(5):
            for k, j in enumerate(range(-1, 3, 2)):
                nlist[i, k] = (i + j) % 5
        nlist = tf.constant(nlist)
        # 5 nodes, 2 neighbors, 4 feature
        edges = tf.ones((5, 2))
        # make inv degree (each one has two connections)
        inv_degree = np.ones((5,)) / 2
        inputs = [nodes, nlist, edges, inv_degree]

        # make peak standards
        ps = {}
        for i in range(16):
            ps[i] = ('F', 0, 1)

        model = nmrgnn.build_GNNModel()
        out_nodes = model(inputs)
        # one peak per atom
        assert out_nodes.shape == (nodes.shape[0],)

        out_nodes = model(inputs)
        assert out_nodes.shape == (nodes.shape[0],)
示例#2
0
    def test_model(self):
        nodes = tf.one_hot([2, 4, 1, 3, 3], 16)
        # neighbors -> 5 nodes, 2 neighbors
        nlist = np.zeros((5, 2), dtype=np.int)
        # make each atom neighbors with subsequent 2 atoms, mod 5
        for i in range(5):
            for k, j in enumerate(range(-1, 3, 2)):
                nlist[i, k] = (i + j) % 5
        nlist = tf.constant(nlist)
        # 5 nodes, 2 neighbors, 4 feature
        edges = tf.ones((5, 2))
        # make inv degree (each one has two connections)
        inv_degree = np.ones((5,)) / 2
        inputs = [nodes, nlist, edges, inv_degree]

        # make peak standards
        ps = {}
        for i in range(16):
            ps[i] = ('F', 0, 1)

        model = nmrgnn.build_GNNModel()
        out_nodes = model(inputs)
        try:
            model.save('gnn_model_load_test')
            del model
            tf.keras.models.load_model(
                'gnn_model_load_test', custom_objects=nmrgnn.custom_objects)
        finally:
            shutil.rmtree('gnn_model_load_test')
示例#3
0
 def test_edgeFCBlock_call(self):
     model = nmrgnn.build_GNNModel()
     hypers = model.hypers
     edge_input = tf.ones((5, 2, 2))
     edgeFCBlock = nmrgnn.EdgeFCBlock(hypers)
     edge_output = edgeFCBlock(edge_input)
     assert edge_output.shape[-1] == hypers.get('edge_feature_size')
     assert edge_output.shape[:-1] == edge_input.shape[:-1]
示例#4
0
 def test_fcBlock_call(self):
     model = nmrgnn.build_GNNModel()
     hypers = model.hypers
     nodes = tf.ones((5, hypers.get('atom_feature_size')))
     fcBlock = nmrgnn.FCBlock(hypers)
     new_nodes = fcBlock(nodes)
     assert new_nodes.shape[-1] == hypers.get('atom_feature_size') // 2
     assert new_nodes.shape[-1] == nodes.shape[-1] // 2
示例#5
0
def train(tfrecords, name, epochs, embeddings, validation, checkpoint_path,
          tensorboard, load, loss_balance):
    '''Train the model'''

    model = nmrgnn.build_GNNModel(loss_balance=loss_balance)
    if load:
        model.load_weights(checkpoint_path)
    callbacks = []
    # set-up learning rate scheduler
    reduce_lr = tf.keras.callbacks.ReduceLROnPlateau(monitor='val_loss',
                                                     factor=0.99,
                                                     patience=4,
                                                     min_lr=1e-4,
                                                     verbose=1)
    callbacks.append(reduce_lr)
    # tensorboard
    if tensorboard is not None:
        tensorboard_callback = tf.keras.callbacks.TensorBoard(
            log_dir=tensorboard,
            write_images=False,
            write_graph=False,
            histogram_freq=0,
            profile_batch=0)
        callbacks.append(tensorboard_callback)
    # save model
    model_checkpoint_callback = tf.keras.callbacks.ModelCheckpoint(
        filepath=checkpoint_path,
        save_weights_only=True,
        monitor='val_loss',
        save_best_only=False)
    callbacks.append(model_checkpoint_callback)

    train_data, validation_data = load_data(tfrecords,
                                            validation,
                                            embeddings,
                                            scale=False)

    # explicitly call model to get shapes defined
    for t in train_data:
        x, y, m = t
        model(x)
        break

    results = model.fit(train_data,
                        epochs=epochs,
                        callbacks=callbacks,
                        validation_data=validation_data,
                        validation_freq=1)

    model.save(name)

    pfile = name + '-history-0.pb'
    i = 0
    while os.path.exists(pfile):
        i += 1
        pfile = f'{name}-history-{i}.pb'
    with open(pfile, 'wb') as f:
        pickle.dump(results.history, file=f)
示例#6
0
 def test_mpBlock_call(self):
     model = nmrgnn.build_GNNModel()
     hypers = model.hypers
     nodes = tf.one_hot([2, 0, 1, 3, 3], 16)
     # neighbors ->  5 nodes, 2 neighbors
     nlist = np.zeros((5, 2), dtype=np.int)
     # make each atom neighbors with subsequent 2 atoms, mod 5
     for i in range(5):
         for k, j in enumerate(range(-1, 3, 2)):
             nlist[i, k] = (i + j) % 5
     nlist = tf.constant(nlist)
     # 5 nodes, 2 neighbors, 4 feature
     edges = tf.ones((5, 2, 2))
     # make inv degree (each one has two connections)
     inv_degree = np.ones((5,)) / 2
     # shapes are specified inside the block
     mp_block = nmrgnn.MPBlock(hypers)
     out_nodes = mp_block([nodes, nlist, edges, inv_degree])
     assert out_nodes.shape == nodes.shape