Exemplo n.º 1
0
    def build_graph(self):
        """Building graph structures:
            Features => DTNNEmbedding => DTNNStep => DTNNStep => DTNNGather => Regression
            """
        self.atom_number = Feature(shape=(None, ), dtype=tf.int32)
        self.distance = Feature(shape=(None, self.n_distance))
        self.atom_membership = Feature(shape=(None, ), dtype=tf.int32)
        self.distance_membership_i = Feature(shape=(None, ), dtype=tf.int32)
        self.distance_membership_j = Feature(shape=(None, ), dtype=tf.int32)

        dtnn_embedding = DTNNEmbedding(n_embedding=self.n_embedding,
                                       in_layers=[self.atom_number])
        if self.dropout > 0.0:
            dtnn_embedding = Dropout(self.dropout, in_layers=dtnn_embedding)
        dtnn_layer1 = DTNNStep(n_embedding=self.n_embedding,
                               n_distance=self.n_distance,
                               in_layers=[
                                   dtnn_embedding, self.distance,
                                   self.distance_membership_i,
                                   self.distance_membership_j
                               ])
        if self.dropout > 0.0:
            dtnn_layer1 = Dropout(self.dropout, in_layers=dtnn_layer1)
        dtnn_layer2 = DTNNStep(n_embedding=self.n_embedding,
                               n_distance=self.n_distance,
                               in_layers=[
                                   dtnn_layer1, self.distance,
                                   self.distance_membership_i,
                                   self.distance_membership_j
                               ])
        if self.dropout > 0.0:
            dtnn_layer2 = Dropout(self.dropout, in_layers=dtnn_layer2)
        dtnn_gather = DTNNGather(n_embedding=self.n_embedding,
                                 layer_sizes=[self.n_hidden],
                                 n_outputs=self.n_tasks,
                                 output_activation=self.output_activation,
                                 in_layers=[dtnn_layer2, self.atom_membership])
        if self.dropout > 0.0:
            dtnn_gather = Dropout(self.dropout, in_layers=dtnn_gather)

        n_tasks = self.n_tasks
        weights = Weights(shape=(None, n_tasks))
        labels = Label(shape=(None, n_tasks))
        output = Reshape(
            shape=(None, n_tasks),
            in_layers=[Dense(in_layers=dtnn_gather, out_channels=n_tasks)])
        self.add_output(output)
        weighted_loss = ReduceSum(L2Loss(in_layers=[labels, output, weights]))
        self.set_loss(weighted_loss)
Exemplo n.º 2
0
  def build_graph(self):
    """Building graph structures:
    Features => DTNNEmbedding => DTNNStep => DTNNStep => DTNNGather => Regression
    """
    self.atom_number = Feature(shape=(None,), dtype=tf.int32)
    self.distance = Feature(shape=(None, self.n_distance))
    self.atom_membership = Feature(shape=(None,), dtype=tf.int32)
    self.distance_membership_i = Feature(shape=(None,), dtype=tf.int32)
    self.distance_membership_j = Feature(shape=(None,), dtype=tf.int32)

    dtnn_embedding = DTNNEmbedding(
        n_embedding=self.n_embedding, in_layers=[self.atom_number])
    dtnn_layer1 = DTNNStep(
        n_embedding=self.n_embedding,
        n_distance=self.n_distance,
        in_layers=[
            dtnn_embedding, self.distance, self.distance_membership_i,
            self.distance_membership_j
        ])
    dtnn_layer2 = DTNNStep(
        n_embedding=self.n_embedding,
        n_distance=self.n_distance,
        in_layers=[
            dtnn_layer1, self.distance, self.distance_membership_i,
            self.distance_membership_j
        ])
    dtnn_gather = DTNNGather(
        n_embedding=self.n_embedding,
        layer_sizes=[self.n_hidden],
        n_outputs=self.n_tasks,
        output_activation=self.output_activation,
        in_layers=[dtnn_layer2, self.atom_membership])

    costs = []
    self.labels_fd = []
    for task in range(self.n_tasks):
      regression = DTNNExtract(task, in_layers=[dtnn_gather])
      self.add_output(regression)
      label = Label(shape=(None, 1))
      self.labels_fd.append(label)
      cost = L2Loss(in_layers=[label, regression])
      costs.append(cost)

    all_cost = Stack(in_layers=costs, axis=1)
    self.weights = Weights(shape=(None, self.n_tasks))
    loss = WeightedError(in_layers=[all_cost, self.weights])
    self.set_loss(loss)
Exemplo n.º 3
0
def test_DTNNStep_pickle():
  tg = TensorGraph()
  atom_features = Feature(shape=(None, 30))
  distance = Feature(shape=(None, 100))
  distance_membership_i = Feature(shape=(None,), dtype=tf.int32)
  distance_membership_j = Feature(shape=(None,), dtype=tf.int32)
  DTNN = DTNNStep(in_layers=[
      atom_features, distance, distance_membership_i, distance_membership_j
  ])
  tg.add_output(DTNN)
  tg.set_loss(DTNN)
  tg.build()
  tg.save()