| 375 | |
| 376 | |
| 377 | def generate_data_tensor(X, Y, batch_size, shuffle=True, num_threads=1, |
| 378 | capacity=None): |
| 379 | #TODO: Add a way with no batch? |
| 380 | #TODO: Set threads to #CPUs fo machine |
| 381 | cr = None |
| 382 | if capacity is None: |
| 383 | capacity = batch_size * num_threads * 4 |
| 384 | |
| 385 | if isinstance(X, tf.Tensor) and isinstance(Y, tf.Tensor): |
| 386 | # Optional Image and Label Batching |
| 387 | if shuffle: |
| 388 | X, Y = tf.train.shuffle_batch([X, Y], batch_size=batch_size, |
| 389 | min_after_dequeue=batch_size, |
| 390 | capacity=capacity, |
| 391 | num_threads=num_threads) |
| 392 | else: |
| 393 | X, Y = tf.train.batch([X, Y], batch_size=batch_size, |
| 394 | capacity=capacity, |
| 395 | num_threads=num_threads) |
| 396 | |
| 397 | # Array Input |
| 398 | elif X is not None and Y is not None: |
| 399 | X_shape = list(np.shape(X)) |
| 400 | Y_shape = list(np.shape(Y)) |
| 401 | # Create a queue using feed_dicts |
| 402 | cr = ArrayFlow(X, Y, batch_size=batch_size, shuffle=shuffle, |
| 403 | capacity=capacity) |
| 404 | X, Y = cr.get() |
| 405 | # Assign a shape to tensors |
| 406 | X_reshape = [-1] + X_shape[1:] if len(X_shape[1:]) > 0 else [-1, 1] |
| 407 | Y_reshape = [-1] + Y_shape[1:] if len(Y_shape[1:]) > 0 else [-1, 1] |
| 408 | X = tf.reshape(X, X_reshape) |
| 409 | Y = tf.reshape(Y, Y_reshape) |
| 410 | |
| 411 | return X, Y, cr |