common.py 6.5 KB
Newer Older
D
dangqingqing 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14
# Copyright (c) 2016 PaddlePaddle Authors. 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.

15
import requests
Y
Yu Yang 已提交
16
import hashlib
Y
Yu Yang 已提交
17
import os
Y
Yu Yang 已提交
18
import shutil
H
Helin Wang 已提交
19
import sys
20 21
import importlib
import paddle.v2.dataset
22 23
import cPickle
import glob
Y
Yu Yang 已提交
24

25
__all__ = ['DATA_HOME', 'download', 'md5file', 'split', 'cluster_files_reader']
Y
Yu Yang 已提交
26

27
DATA_HOME = os.path.expanduser('~/.cache/paddle/dataset')
Y
Yu Yang 已提交
28 29 30

if not os.path.exists(DATA_HOME):
    os.makedirs(DATA_HOME)
Y
Yu Yang 已提交
31 32


33 34 35 36 37 38 39
def md5file(fname):
    hash_md5 = hashlib.md5()
    f = open(fname, "rb")
    for chunk in iter(lambda: f.read(4096), b""):
        hash_md5.update(chunk)
    f.close()
    return hash_md5.hexdigest()
40 41 42 43 44 45 46 47 48


def download(url, module_name, md5sum):
    dirname = os.path.join(DATA_HOME, module_name)
    if not os.path.exists(dirname):
        os.makedirs(dirname)

    filename = os.path.join(dirname, url.split('/')[-1])
    if not (os.path.exists(filename) and md5file(filename) == md5sum):
H
Helin Wang 已提交
49
        print "Cache file %s not found, downloading %s" % (filename, url)
50
        r = requests.get(url, stream=True)
H
Helin Wang 已提交
51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66
        total_length = r.headers.get('content-length')

        if total_length is None:
            with open(filename, 'w') as f:
                shutil.copyfileobj(r.raw, f)
        else:
            with open(filename, 'w') as f:
                dl = 0
                total_length = int(total_length)
                for data in r.iter_content(chunk_size=4096):
                    dl += len(data)
                    f.write(data)
                    done = int(50 * dl / total_length)
                    sys.stdout.write("\r[%s%s]" % ('=' * done,
                                                   ' ' * (50 - done)))
                    sys.stdout.flush()
67 68

    return filename
Y
Yi Wang 已提交
69 70


71 72 73 74 75 76 77 78
def fetch_all():
    for module_name in filter(lambda x: not x.startswith("__"),
                              dir(paddle.v2.dataset)):
        if "fetch" in dir(
                importlib.import_module("paddle.v2.dataset.%s" % module_name)):
            getattr(
                importlib.import_module("paddle.v2.dataset.%s" % module_name),
                "fetch")()
79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151


def split(reader, line_count, suffix="%05d.pickle", dumper=cPickle.dump):
    """
    you can call the function as:

    split(paddle.v2.dataset.cifar.train10(), line_count=1000,
        suffix="imikolov-train-%05d.pickle")

    the output files as:

    |-imikolov-train-00000.pickle
    |-imikolov-train-00001.pickle
    |- ...
    |-imikolov-train-00480.pickle

    :param reader: is a reader creator
    :param line_count: line count for each file
    :param suffix: the suffix for the output files, should contain "%d"
                means the id for each file. Default is "%05d.pickle"
    :param dumper: is a callable function that dump object to file, this
                function will be called as dumper(obj, f) and obj is the object
                will be dumped, f is a file object. Default is cPickle.dump.
    """
    if not callable(dumper):
        raise TypeError("dumper should be callable.")
    lines = []
    indx_f = 0
    for i, d in enumerate(reader()):
        lines.append(d)
        if i >= line_count and i % line_count == 0:
            with open(suffix % indx_f, "w") as f:
                dumper(lines, f)
                lines = []
                indx_f += 1
    if lines:
        with open(suffix % indx_f, "w") as f:
            dumper(lines, f)


def cluster_files_reader(files_pattern,
                         trainer_count,
                         trainer_id,
                         loader=cPickle.load):
    """
    Create a reader that yield element from the given files, select
    a file set according trainer count and trainer_id

    :param files_pattern: the files which generating by split(...)
    :param trainer_count: total trainer count
    :param trainer_id: the trainer rank id
    :param loader: is a callable function that load object from file, this
                function will be called as loader(f) and f is a file object.
                Default is cPickle.load
    """

    def reader():
        if not callable(loader):
            raise TypeError("loader should be callable.")
        file_list = glob.glob(files_pattern)
        file_list.sort()
        my_file_list = []
        for idx, fn in enumerate(file_list):
            if idx % trainer_count == trainer_id:
                print "append file: %s" % fn
                my_file_list.append(fn)
        for fn in my_file_list:
            with open(fn, "r") as f:
                lines = loader(f)
                for line in lines:
                    yield line

    return reader
G
gongweibao 已提交
152 153


G
gongweibao 已提交
154 155 156 157
def convert(output_path,
            reader,
            num_shards,
            name_prefix,
G
gongweibao 已提交
158
            max_lines_to_shuffle=1000):
G
gongweibao 已提交
159 160
    import recordio
    import cPickle as pickle
G
gongweibao 已提交
161
    import random
G
gongweibao 已提交
162 163 164 165 166 167 168
    """
    Convert data from reader to recordio format files.

    :param output_path: directory in which output files will be saved.
    :param reader: a data reader, from which the convert program will read data instances.
    :param num_shards: the number of shards that the dataset will be partitioned into.
    :param name_prefix: the name prefix of generated files.
G
gongweibao 已提交
169
    :param max_lines_to_shuffle: the max lines numbers to shuffle before writing.
G
gongweibao 已提交
170 171
    """

G
gongweibao 已提交
172 173
    assert num_shards >= 1
    assert max_lines_to_shuffle >= 1
G
gongweibao 已提交
174

G
gongweibao 已提交
175 176 177 178 179 180
    def open_writers():
        w = []
        for i in range(0, num_shards):
            n = "%s/%s-%05d-of-%05d" % (output_path, name_prefix, i,
                                        num_shards - 1)
            w.append(recordio.writer(n))
G
gongweibao 已提交
181

G
gongweibao 已提交
182
        return w
G
gongweibao 已提交
183

G
gongweibao 已提交
184 185 186
    def close_writers(w):
        for i in range(0, num_shards):
            w[i].close()
G
gongweibao 已提交
187

G
gongweibao 已提交
188 189 190 191 192
    def write_data(w, lines):
        random.shuffle(lines)
        for i, d in enumerate(lines):
            d = pickle.dumps(d, pickle.HIGHEST_PROTOCOL)
            w[i % num_shards].write(d)
G
gongweibao 已提交
193

G
gongweibao 已提交
194 195
    w = open_writers()
    lines = []
G
gongweibao 已提交
196

G
gongweibao 已提交
197 198 199 200 201 202
    for i, d in enumerate(reader()):
        lines.append(d)
        if i % max_lines_to_shuffle == 0 and i >= max_lines_to_shuffle:
            write_data(w, lines)
            lines = []
            continue
G
gongweibao 已提交
203

G
gongweibao 已提交
204 205
    write_data(w, lines)
    close_writers(w)