这是pytorch中用于加载数据的两个类,在torch.utils.data包下
。
首先继承Dataset类,实现三个方法:__init__(),__getitem__(),__len__()
。init方法中初始化数据,getitem方法中根据指定的index找到数据,len方法返回数据的长度。
之后我们new一个Dataloder类,传入我们new的Dataset对象,指定每次需要取得数据的batch_size,shuffle即数据是否要打乱。new完这个对象之后,就可以遍历获得数据。
import pandas as pd
import torch
from torch.utils.data import Dataset, DataLoader
from torch.utils.data.dataset import T_co
# 将实体和关系映射到id上,方便处理
entity_id=pd.read_csv('./data/entity2id.csv')
entity_to_id={
entity_id.iloc[i]['name']:entity_id.iloc[i