Python utils.safe_izip函数代码示例

def write_h5(full_dataset, args, rng):
    partition_masks = get_partition_masks(full_dataset, args.ratio, rng)
    partition_sizes = [numpy.count_nonzero(mask) for mask in partition_masks]

    basename = os.path.splitext(os.path.split(args.input)[1])[0]

    output_path = os.path.join(args.output_dir,
                               "{}_split_{}.h5".format(basename, args.ratio))

    h5_file = make_h5_file(output_path,
                           args.partition_names,
                           partition_sizes,
                           full_dataset.names,
                           full_dataset.formats)

    for partition_name, partition_mask in safe_izip(args.partition_names,
                                                    partition_masks):
        partition = h5_file['partitions'][partition_name]

        for full_tensor, name, fmt in safe_izip(full_dataset.tensors,
                                                full_dataset.names,
                                                full_dataset.formats):
            partition_tensor = partition[name]
            batch_slice = get_batch_slice(fmt, partition_mask)
            partition_tensor[...] = full_tensor[batch_slice]

def write_memmaps(full_dataset, args, rng):

    def get_partition_path(output_dir, input_path, partition_name):
        basename = os.path.splitext(os.path.split(input_path)[1])[0]
        return os.path.join(output_dir,
                            "{}_split_{}_{}.npy".format(basename,
                                                        args.ratio,
                                                        partition_name))

    partition_masks = get_partition_masks(full_dataset, args.ratio, rng)

    for partition_name, partition_mask in safe_izip(args.partition_names,
                                                    partition_masks):
        partition_path = get_partition_path(args.output_dir,
                                            args.input,
                                            partition_name)

        memmap = make_memmap_file(partition_path,
                                  numpy.count_nonzero(partition_mask),
                                  full_dataset.names,
                                  full_dataset.formats)


        for full_tensor, name, fmt in safe_izip(full_dataset.tensors,
                                                full_dataset.names,
                                                full_dataset.formats):
            partition_tensor = memmap[name]
            batch_slice = get_batch_slice(fmt, partition_mask)
            partition_tensor[...] = full_tensor[batch_slice]

def main():
    args = parse_args()

    h5_datasets = load_h5_dataset(args.input)

    def get_output_filepath(input_name, partition_name):
        dirname, filename = os.path.split(input_name)
        basename, extension = os.path.splitext(filename)
        assert_equal(extension, '.h5')

        return os.path.join(dirname, "{}_{}{}".format(basename,
                                                      partition_name,
                                                      '.npy'))

    for dataset, partition_name in safe_izip(h5_datasets, h5_datasets._fields):
        output_filepath = get_output_filepath(args.input, partition_name)
        memmap = make_memmap_file(output_filepath,
                                  dataset.num_examples(),
                                  dataset.names,
                                  dataset.formats)

        memmap_tensors = [memmap[name] for name in dataset.names]
        for in_tensor, out_tensor in safe_izip(dataset.tensors,
                                               memmap_tensors):
            assert_equal(out_tensor.shape, in_tensor.shape)
            out_tensor[...] = in_tensor


        print("Wrote {}".format(output_filepath))

def test_mnist():
    '''
    Tests load_mnist().

    Checks test & train sets' formats and sizes, but not content.
    '''
    train_set, test_set = load_mnist()

    for mnist, expected_size in safe_izip((train_set, test_set),
                                          (60000, 10000)):
        assert_equal(mnist.num_examples(), expected_size)

    expected_formats = [DenseFormat(shape=[-1, 28, 28],
                                    axes=['b', '0', '1'],
                                    dtype='uint8'),
                        DenseFormat(shape=[-1],
                                    axes=['b'],
                                    dtype='uint8')]
    expected_names = [u'images', u'labels']
    expected_sizes = [60000, 10000]

    for dataset, expected_size in safe_izip((train_set, test_set),
                                            expected_sizes):
        assert_all_equal(dataset.names, expected_names)
        assert_all_equal(dataset.formats, expected_formats)

        for tensor, fmt in safe_izip(dataset.tensors, dataset.formats):
            fmt.check(tensor)
            assert_equal(tensor.shape[0], expected_size)

        labels = dataset.tensors[dataset.names.index('labels')]
        assert_true(numpy.logical_and(labels[...] >= 0,
                                      labels[...] < 10).all())

def test_pickle_h5_dataset():
    '''
    Tests pickling and unpickling of H5Dataset.
    '''

    # Path for the pickle file, not the .h5 file.
    file_path = '/tmp/test_mnist_test_pickle_hdf5_data.pkl'

    def make_pickle(file_path):
        '''
        Pickles the MNIST dataset.
        '''
        hdf5_data = load_mnist()
        with open(file_path, 'wb') as pickle_file:
            cPickle.dump(hdf5_data,
                         pickle_file,
                         protocol=cPickle.HIGHEST_PROTOCOL)

    make_pickle(file_path)
    assert_less(_file_size_in_bytes(file_path), 1024 * 5)

    def load_pickle(file_path):
        '''
        Loads the MNIST dataset pickled above.
        '''
        with open(file_path, 'rb') as pickle_file:
            return cPickle.load(pickle_file)

    mnist_datasets = load_mnist()
    pickled_mnist_datasets = load_pickle(file_path)

    for (mnist_dataset,
         pickled_mnist_dataset) in safe_izip(mnist_datasets,
                                             pickled_mnist_datasets):
        for (name,
             expected_name,
             fmt,
             expected_fmt,
             tensor,
             expected_tensor) in safe_izip(pickled_mnist_dataset.names,
                                           mnist_dataset.names,
                                           pickled_mnist_dataset.formats,
                                           mnist_dataset.formats,
                                           pickled_mnist_dataset.tensors,
                                           mnist_dataset.tensors):
            assert_equal(name, expected_name)
            assert_equal(fmt, expected_fmt)
            assert_array_equal(tensor, expected_tensor)

    def draw_menu():
        lines = []

        rows, rows_index = label_index_map.index_to_rows(index)

        if rows is not None:
            row = rows[rows_index[0]]
            label = dataset_labels[row]

            for name, value, converter in safe_izip(label_names[:5],
                                                    label[:5],
                                                    converters[:5]):
                lines.append('{}: {}'.format(name, converter(value)))

            lines.append("image: {} of {}".format(rows_index[0] + 1,
                                                  len(rows)))
            lines.append('')

            if dataset_labels.shape[1] == 11:
                for name, value, converter in safe_izip(label_names[5:],
                                                        label[5:],
                                                        converters[5:]):
                    lines.append('{}: {}'.format(name, converter(value)))
        else:
            label_5d = label_index_map.index_to_label_5d(index)

            for name, value, converter in safe_izip(label_names[:5],
                                                    label_5d,
                                                    converters[:5]):
                lines.append('{}: {}'.format(name, converter(value)))

            lines.append("image: (no such image)")
            lines.append('')

            if dataset_labels.shape[1] == 11:
                for name in label_names[5:]:
                    lines.append('{}: N/A'.format(name))

        lines[index_dim[0]] = "==> " + lines[index_dim[0]]

        # "transAxes": 0, 0 = bottom-left, 1, 1 at upper-right.
        text_axes = all_axes[0]
        text_axes.clear()
        text_axes.text(0.1,  # x
                       0.5,  # y
                       '\n'.join(lines),
                       verticalalignment='center',
                       transform=text_axes.transAxes)

    def __getitem__(self, arg):
        '''
        Returns a tuple of slices along the batch axis.

        Examples:

          self[slice(3, 10)] returns a 6-element batch from each tensor.

          self[3] is equivalent to self[slice(3, 4)] (batch axis is intact)


        Parameters
        ----------
        arg: integer or slice

        Returns
        -------
        rval: tuple
          A tuple of views into each tensor.
        '''

        batch_slice = self._getitem_arg_to_slice(arg)

        def get_slice_tuple(fmt, batch_slice):
            '''
            Returns a tuple for slicing a tensor along its batch axis.
            '''
            return tuple(batch_slice if axis == 'b'
                         else slice(None)
                         for axis in fmt.axes)

        return tuple(tensor[get_slice_tuple(fmt, batch_slice)]
                     for tensor, fmt
                     in safe_izip(self.tensors, self.formats))

    def draw_images():
        def draw_image_impl(image, axes):
            if args.lcn:
                axes.imshow(image, cmap='gray')
            else:
                axes.imshow(image,
                            cmap='gray',
                            norm=matplotlib.colors.NoNorm())


        rows, rows_index = label_index_map.index_to_rows(index)

        image_axes = all_axes[1:]  # could be 1 or two axes

        if rows is None:
            for axes in image_axes:
                axes.clear()
        else:
            row = rows[rows_index[0]]
            image = dataset_images[row]
            if args.lcn:
                image = lcn(image)

            if 's' in image_format.axes:
                assert_equal(image_format.axes.index('s'), 1)
                for sub_image, axes in safe_izip(image, image_axes):
                    draw_image_impl(sub_image, axes)
            else:
                draw_image_impl(image, all_axes[1])

    def __init__(self, names, tensors, formats):
        if len(names) != len(formats) or len(names) != len(tensors):
            raise ValueError("Names, formats, and tensors must all have the "
                             "same length, but got %d, %d, and %d "
                             "respectively." %
                             tuple(len(names), len(formats), len(tensors)))

        for name in names:
            if not isinstance(name, basestring):
                raise TypeError("names must be strings, not %s." % type(name))

        for tensor, fmt in safe_izip(tensors, formats):
            if not isinstance(fmt, Format):
                raise TypeError("formats must be Formats, not %s." %
                                type(fmt))

            if 'b' not in fmt.axes:
                raise ValueError("Expected format to contain a 'b' axis "
                                 "(batch axis).")

            if fmt.dtype is None:
                raise ValueError("Expected all formats to specify a dtype.")

            fmt.check(tensor)

        self.names = tuple(names)
        self.formats = tuple(formats)
        self.tensors = tuple(tensors)

    def label5d_to_index(self, labels):
        assert_in(len(labels.shape), (1, 2))

        was_1D = False
        if len(labels.shape) == 1:
            labels = labels[numpy.newaxis, :]
            was_1D = True

        labels = numpy.asarray(labels[:, :5])

        indices = numpy.empty_like(labels)
        indices[...] = -1

        for (label5d_values,
             label_column,
             index_column) in safe_izip(self.label5d_values,
                                        labels.T,
                                        indices.T):

            for ind, value in enumerate(label5d_values):
                mask = (label_column == value)
                index_column[mask] = ind

        assert_false((indices == -1).any())

        if was_1D:
            assert_equal(indices.shape[0], 1)
            return indices[0, :]
        else:
            return indices

    def _next(self):

        batch_indices = self._next_batch_indices()
        # pdb.set_trace()

        # sanity-check output of _next_batch_indices()
        if not isinstance(batch_indices, slice):
            assert_all_integer(batch_indices)

            if isinstance(batch_indices, numpy.ndarray):
                # Workaround to a bug in h5py.Dataset where indexing by a
                # length-1 ndarray is treated like indexing with the integer it
                # contains.
                if len(batch_indices) == 1:
                    batch_indices = tuple(batch_indices)
            else:
                assert_is_instance(batch_indices, collections.Sequence)

        result = tuple(self._get_batches(self.dataset.tensors,
                                         self.dataset.formats,
                                         batch_indices))

        # sanity-check size of batches
        for batch, fmt in safe_izip(result, self.dataset.formats):
            assert_equal(batch.shape[fmt.axes.index('b')], self.batch_size)

        return result

    def _get_batches(self, tensors, formats, batch_indices):
        '''
        Extracts batches from tensors, given batch_indices.

        Parameters
        ----------
        tensors: Iterable of numpy.ndarray, or similar
          The tensors to select a batch from. Usually self.dataset.tensors.

        fmt: simplelearn.format.DenseFormat
          The formats corresponding to <tensors>. Usually self.dataset.formats.

        batch_indices: Sequence
          The output of _get_batch_indices.
        '''

        def get_batch(tensor, fmt):

            # h5py has a bug where if len(index_tuple) == 1,
            # Dataset.__getitem__ treats it the same as just
            # index_tuple[0]. Adding a gratuitous Ellipsis element to the end
            # prevents this.
            #
            # See h5py bug report: https://github.com/h5py/h5py/issues/586
            index_tuple = tuple(batch_indices if axis == 'b' else slice(None)
                                for axis in fmt.axes) + (Ellipsis, )

            return tensor[index_tuple]

        return tuple(get_batch(tensor, fmt) for tensor, fmt
                     in safe_izip(tensors, formats))

    def index_to_label_5d(self, indices):
        assert_in(len(indices.shape), (1, 2))

        indices = numpy.asarray(indices)

        was_1D = False
        if len(indices.shape) == 1:
            indices = indices[numpy.newaxis, :]
            was_1D = True

        assert_equal(indices.shape[1], 5)

        labels = numpy.zeros_like(indices)

        for (index_column,
             label_column,
             label5d_values) in safe_izip(indices.T,
                                          labels.T,
                                          self.label5d_values):
            label_column[:] = label5d_values[index_column]

        if was_1D:
            return labels[0, :]
        else:
            return labels

    def __init__(self, path):
        assert_is_instance(path, basestring)

        path = os.path.abspath(path)

        assert_true(path.startswith(simplelearn.data.data_path),
                    ("{} is not a subdirectory of simplelearn.data.data_path "
                     "{}").format(path, simplelearn.data.data_path))

        # pylint can't see memmap members
        # pylint: disable=no-member
        self.memmap = numpy.lib.format.open_memmap(path, mode='r')
        num_examples = self.memmap.shape[0]

        names = self.memmap.dtype.names
        tensors = [self.memmap[name] for name in names]
        axes_list = [field[2] for field
                     in self.memmap.dtype.fields.itervalues()]

        def replace_element(arg, index, new_value):
            assert_is_instance(arg, tuple)
            result = list(arg)
            result[index] = new_value
            return tuple(result)

        formats = [DenseFormat(axes=axes,
                               shape=replace_element(tensor.shape, 0, -1),
                               dtype=tensor.dtype)
                   for axes, tensor in safe_izip(axes_list, tensors)]

        super(MemmapDataset, self).__init__(names=names,
                                            formats=formats,
                                            tensors=tensors)

        def make_id_elev_azim_to_example(example_dataset):
            assert_equal(example_dataset.formats[0].axes.index('c'), 3)
            assert_equal(example_dataset.formats[0].axes.index('b'), 0)
            assert_equal(example_dataset.formats[1].axes.index('b'), 0)

            images = example_dataset.tensors[0][..., :3]  # cut off alpha
            labels = example_dataset.tensors[1]

            assert_in(labels.shape[1], (5, 11))

            # Arbitrarily restrict attention to images that use
            # lighting setup 0.
            row_mask = labels[:, 4] == 0
            images = images[row_mask, :]
            labels = labels[row_mask, :]

            ids = self.label_to_id(labels)
            ids_elevs_azims = numpy.hstack((ids[:, numpy.newaxis],
                                            labels[:, 2:4]))

            result = dict(safe_izip((tuple(t) for t in ids_elevs_azims),
                                    images))

            assert_equal(len(result), ids_elevs_azims.shape[0])

            return result