用pytorch来制作数据集(分为有y和无y两种)

def get_data(path,batch_size,shuffle=True,y_value=False):
    df = pd.read_csv(path,index_col=0)
    if(y_value):
        x = np.array(df.iloc[:,:-1])
        y =  np.array(df.iloc[:,-1])
        train_dataset = TensorDataset(torch.from_numpy(x).unsqueeze(1).float(),torch.from_numpy(y))
        train_iter = DataLoader(dataset=train_dataset,batch_size=batch_size,shuffle = shuffle)
    else:
        x = np.array(df)
        train_dataset = TensorDataset(torch.tensor(x).unsqueeze(1).float())
        train_iter = DataLoader(dataset=train_dataset,batch_size=batch_size, shuffle=shuffle)
    return train_iter


for x in data:        #此处的data是由get_data生成的
    if(len(x) == 2):  #因为生成的数据可能有y和无y,首先通过判断长度来确定
        data_x = x[0]
        data_y = x[1]
        print(data_x.shape)
        print(data_y.shape)
    else:
        data_x = x[0]