Spaces:
Running
Running
File size: 2,576 Bytes
0b8359d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 |
# 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)
|