Spaces:
Running
Running
# Copyright 2018 Google, Inc. All Rights Reserved. | |
# | |
# Licensed under the Apache License, Version 2.0 (the "License"); | |
# you may not use this file except in compliance with the License. | |
# You may obtain a copy of the License at | |
# | |
# http://www.apache.org/licenses/LICENSE-2.0 | |
# | |
# Unless required by applicable law or agreed to in writing, software | |
# distributed under the License is distributed on an "AS IS" BASIS, | |
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
# See the License for the specific language governing permissions and | |
# limitations under the License. | |
# ============================================================================== | |
import sonnet as snt | |
import tensorflow as tf | |
from tensorflow.python.keras.datasets import mnist | |
from learning_unsupervised_learning.datasets import common | |
class Mnist(snt.AbstractModule): | |
def __init__(self, device, batch_size=128, name="Mnist"): | |
self.device = device | |
self.batch_size = batch_size | |
self._make_dataset() | |
self.iterator = None | |
super(Mnist, self).__init__(name=name) | |
def _make_dataset(self): | |
(x_train, y_train), (x_test, y_test) = mnist.load_data() | |
x_train = x_train.reshape(60000, 784) | |
x_test = x_test.reshape(10000, 784) | |
dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train)) | |
dataset = dataset.repeat() | |
dataset = dataset.shuffle(self.batch_size * 3) | |
dataset = dataset.batch(self.batch_size) | |
def _map_fn(image, label): | |
image = tf.to_float(image) / 255. | |
label.set_shape([self.batch_size]) | |
label = tf.cast(label, dtype=tf.int32) | |
label_onehot = tf.one_hot(label, 10) | |
image = tf.reshape(image, [self.batch_size, 28, 28, 1]) | |
return common.ImageLabelOnehot( | |
image=image, label=label, label_onehot=label_onehot) | |
self.dataset = dataset.map(_map_fn) | |
def _build(self): | |
if self.iterator is None: | |
self.iterator = self.dataset.make_one_shot_iterator() | |
batch = self.iterator.get_next() | |
[b.set_shape([self.batch_size] + b.shape.as_list()[1:]) for b in batch] | |
return batch | |
class TinyMnist(Mnist): | |
def __init__(self, *args, **kwargs): | |
kwargs.setdefault("name", "TinyMnist") | |
super(TinyMnist, self).__init__(*args, **kwargs) | |
def _make_dataset(self): | |
super(TinyMnist, self)._make_dataset() | |
def _map_fn(batch): | |
new_img = tf.image.resize_images(batch.image, [14, 14]) | |
return common.ImageLabelOnehot( | |
image=new_img, label=batch.label, label_onehot=batch.label_onehot) | |
self.dataset = self.dataset.map(_map_fn) | |