corentin/torch_utils
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Clean mnist comment, add Linear layer

Browse source
This commit is contained in:
Corentin Risselin 2020-04-17 10:45:13 +09:00
commit ced13a4351
2 changed files with 20 additions and 11 deletions
Download patch file Download diff file Expand all files Collapse all files

8
dataset/mnist.py
View file

@ -9,13 +9,13 @@ def load_image_file(path: str, magic_number: int = 2051, flatten: bool = False)
"""Load MNIST image file"""
images = []
with open(path, 'rb') as data_file:
header_data = data_file.read(16) ## 4 * int32 = 16 bytes
header_data = data_file.read(16) # 4 * int32 = 16 bytes
data_magic, data_count, rows, cols = struct.unpack('>iiii', header_data)
if data_magic != magic_number:
raise RuntimeError(
f'MNIST image file doesn\'t have correct mmagic number: {data_magic} instead of {magic_number}')
image_data_size = rows * cols
image_chunk = 1000 ## loading by chunk for faster IO
image_chunk = 1000 # loading by chunk for faster IO
for _ in range(data_count // image_chunk):
data = data_file.read(image_data_size * image_chunk)
data = struct.unpack('B' * image_data_size * image_chunk, data)
@ -32,12 +32,12 @@ def load_label_file(path: str, magic_number: int = 2049) -> np.ndarray:
"""Load MNIST label file"""
labels = []
with open(path, 'rb') as data_file:
header_data = data_file.read(8) ## 2 * int32 = 8 bytes
header_data = data_file.read(8) # 2 * int32 = 8 bytes
data_magic, data_count = struct.unpack('>ii', header_data)
if data_magic != magic_number:
raise RuntimeError(
f'MNIST label file doesn\'t have correct mmagic number: {data_magic} instead of {magic_number}')
label_chunk = 1000 ## loading by chunk for faster IO
label_chunk = 1000 # loading by chunk for faster IO
for _ in range(data_count // label_chunk):
data = data_file.read(label_chunk)
data = struct.unpack('B' * label_chunk, data)