Load

In [ ]:

import torch
import torch.utils.data as Data
import struct
from sklearn.model_selection import train_test_split

import torch import torch.utils.data as Data import struct from sklearn.model_selection import train_test_split

In [ ]:

'''
    This file is mainly about loading the data.
'''
def load_image_ubyte(file):
    with open(file, 'rb') as f:
        bin_data = f.read()
    offset = 0
    fmt_header = '>iiii'
    magic, numImgs, numRows, numCols = struct.unpack_from(fmt_header, bin_data, offset)
    
    offset = struct.calcsize(fmt_header)
    fmt_image = '>' + str(numImgs * numRows * numCols) + 'B'
    data = torch.tensor(struct.unpack_from(fmt_image, bin_data, offset)).reshape(numImgs, 1, numRows, numCols).float()
    return data

''' This file is mainly about loading the data. ''' def load_image_ubyte(file): with open(file, 'rb') as f: bin_data = f.read() offset = 0 fmt_header = '>iiii' magic, numImgs, numRows, numCols = struct.unpack_from(fmt_header, bin_data, offset) offset = struct.calcsize(fmt_header) fmt_image = '>' + str(numImgs * numRows * numCols) + 'B' data = torch.tensor(struct.unpack_from(fmt_image, bin_data, offset)).reshape(numImgs, 1, numRows, numCols).float() return data

In [ ]:

def load_label_ubyte(file):
    with open(file, 'rb') as f:
        bin_data = f.read()
    offset = 0
    fmt_header = '>ii'
    magic, numImgs = struct.unpack_from(fmt_header, bin_data, offset)
    
    offset = struct.calcsize(fmt_header)
    fmt_image = '>' + str(numImgs) + 'B'
    data = torch.tensor(struct.unpack_from(fmt_image, bin_data, offset)).long()
    return data

def load_label_ubyte(file): with open(file, 'rb') as f: bin_data = f.read() offset = 0 fmt_header = '>ii' magic, numImgs = struct.unpack_from(fmt_header, bin_data, offset) offset = struct.calcsize(fmt_header) fmt_image = '>' + str(numImgs) + 'B' data = torch.tensor(struct.unpack_from(fmt_image, bin_data, offset)).long() return data

In [ ]:

'''
    Given datasets:
    train-images-idx3-ubyte : train images
    train-labels-idx1-ubyte : train labels
    t10k-images-idx3-ubyte : test images
    t10k-labels-idx1-ubyte : test labels
    Size: train:60000 test:10000
'''

''' Given datasets: train-images-idx3-ubyte : train images train-labels-idx1-ubyte : train labels t10k-images-idx3-ubyte : test images t10k-labels-idx1-ubyte : test labels Size: train:60000 test:10000 '''

In [ ]:

def load(path, batch_size=128,val_rate=0.2):
    # 先从对应的数据集中解析数据，并将结果张量化
    x_train = load_image_ubyte(path + "/train-images-idx3-ubyte")
    y_train = load_label_ubyte(path + "/train-labels-idx1-ubyte")
    x_test = load_image_ubyte(path + "/t10k-images-idx3-ubyte")
    y_test = load_label_ubyte(path + "/t10k-labels-idx1-ubyte")
    test_data = Data.TensorDataset(x_test, y_test)
    test_data = Data.DataLoader(test_data, batch_size=batch_size, shuffle=True,num_workers=4)

    x_train, x_val, y_train, y_val = train_test_split(x_train, y_train, test_size=val_rate, random_state=42,shuffle=True)

    train_data = Data.TensorDataset(x_train, y_train)
    train_data = Data.DataLoader(train_data, batch_size=batch_size, shuffle=True,num_workers=4)

    val_data = Data.TensorDataset(x_val, y_val)
    val_data = Data.DataLoader(val_data, batch_size=batch_size, shuffle=True,num_workers=4)

    '''
    print(x_train.shape[0])
    print(x_val.shape[0])
    print(x_test.shape[0])
    '''
    return train_data,val_data,test_data

def load(path, batch_size=128,val_rate=0.2): # 先从对应的数据集中解析数据，并将结果张量化 x_train = load_image_ubyte(path + "/train-images-idx3-ubyte") y_train = load_label_ubyte(path + "/train-labels-idx1-ubyte") x_test = load_image_ubyte(path + "/t10k-images-idx3-ubyte") y_test = load_label_ubyte(path + "/t10k-labels-idx1-ubyte") test_data = Data.TensorDataset(x_test, y_test) test_data = Data.DataLoader(test_data, batch_size=batch_size, shuffle=True,num_workers=4) x_train, x_val, y_train, y_val = train_test_split(x_train, y_train, test_size=val_rate, random_state=42,shuffle=True) train_data = Data.TensorDataset(x_train, y_train) train_data = Data.DataLoader(train_data, batch_size=batch_size, shuffle=True,num_workers=4) val_data = Data.TensorDataset(x_val, y_val) val_data = Data.DataLoader(val_data, batch_size=batch_size, shuffle=True,num_workers=4) ''' print(x_train.shape[0]) print(x_val.shape[0]) print(x_test.shape[0]) ''' return train_data,val_data,test_data