Build the MNIST model up to where it may be used for inference. Args: inp: input data num_clusters: number of clusters of input features to train. hidden1_units: Size of the first hidden layer. hidden2_units: Size of the second hidden layer. Returns: logits: Output tensor w
(inp, num_clusters, hidden1_units, hidden2_units)
| 119 | |
| 120 | |
| 121 | def inference(inp, num_clusters, hidden1_units, hidden2_units): |
| 122 | """Build the MNIST model up to where it may be used for inference. |
| 123 | |
| 124 | Args: |
| 125 | inp: input data |
| 126 | num_clusters: number of clusters of input features to train. |
| 127 | hidden1_units: Size of the first hidden layer. |
| 128 | hidden2_units: Size of the second hidden layer. |
| 129 | |
| 130 | Returns: |
| 131 | logits: Output tensor with the computed logits. |
| 132 | clustering_loss: Clustering loss. |
| 133 | kmeans_training_op: An op to train the clustering. |
| 134 | """ |
| 135 | # Clustering |
| 136 | kmeans = tf.contrib.factorization.KMeans( |
| 137 | inp, |
| 138 | num_clusters, |
| 139 | distance_metric=tf.contrib.factorization.COSINE_DISTANCE, |
| 140 | # TODO(agarwal): kmeans++ is currently causing crash in dbg mode. |
| 141 | # Enable this after fixing. |
| 142 | # initial_clusters=tf.contrib.factorization.KMEANS_PLUS_PLUS_INIT, |
| 143 | use_mini_batch=True) |
| 144 | |
| 145 | (all_scores, _, clustering_scores, _, kmeans_init, |
| 146 | kmeans_training_op) = kmeans.training_graph() |
| 147 | # Some heuristics to approximately whiten this output. |
| 148 | all_scores = (all_scores[0] - 0.5) * 5 |
| 149 | # Here we avoid passing the gradients from the supervised objective back to |
| 150 | # the clusters by creating a stop_gradient node. |
| 151 | all_scores = tf.stop_gradient(all_scores) |
| 152 | clustering_loss = tf.reduce_sum(clustering_scores[0]) |
| 153 | # Hidden 1 |
| 154 | with tf.name_scope('hidden1'): |
| 155 | weights = tf.Variable( |
| 156 | tf.truncated_normal([num_clusters, hidden1_units], |
| 157 | stddev=1.0 / math.sqrt(float(IMAGE_PIXELS))), |
| 158 | name='weights') |
| 159 | biases = tf.Variable(tf.zeros([hidden1_units]), |
| 160 | name='biases') |
| 161 | hidden1 = tf.nn.relu(tf.matmul(all_scores, weights) + biases) |
| 162 | # Hidden 2 |
| 163 | with tf.name_scope('hidden2'): |
| 164 | weights = tf.Variable( |
| 165 | tf.truncated_normal([hidden1_units, hidden2_units], |
| 166 | stddev=1.0 / math.sqrt(float(hidden1_units))), |
| 167 | name='weights') |
| 168 | biases = tf.Variable(tf.zeros([hidden2_units]), |
| 169 | name='biases') |
| 170 | hidden2 = tf.nn.relu(tf.matmul(hidden1, weights) + biases) |
| 171 | # Linear |
| 172 | with tf.name_scope('softmax_linear'): |
| 173 | weights = tf.Variable( |
| 174 | tf.truncated_normal([hidden2_units, NUM_CLASSES], |
| 175 | stddev=1.0 / math.sqrt(float(hidden2_units))), |
| 176 | name='weights') |
| 177 | biases = tf.Variable(tf.zeros([NUM_CLASSES]), |
| 178 | name='biases') |
no test coverage detected