数据切分训练验证集时出现KeyError:0不在范围内问题求助

错误根源 #

你遇到的KeyError本质是切分后的DataFrame索引未重置，导致Dataset中连续的0、1、2...索引无法匹配DataFrame的非连续索引。比如执行train_df = df[int(size*0.2):]后，train_df的起始索引是int(size*0.2)而非0，但你的__getitem__方法里用index（从0开始）去访问self.df["ddg"][index]，自然找不到对应键。

未切分时正常是因为原始DataFrame的索引是从0开始的连续整数，index和DataFrame索引完全对应。

两种解决方案 #

方案1：切分后重置DataFrame索引 #

在切分DataFrame时添加reset_index(drop=True)，强制重置为从0开始的连续索引：

# 修改数据切分代码
train_wt_emb = wt_emb[int(size*0.2):]
train_mut_emb = mut_emb[int(size*0.2):]
train_df = df[int(size*0.2):].reset_index(drop=True)  # 重置索引

valid_wt_emb = wt_emb[:int(size*0.2)]
valid_mut_emb = mut_emb[:int(size*0.2)]
valid_df = df[:int(size*0.2)].reset_index(drop=True)  # 重置索引

方案2：在Dataset中用位置索引访问DataFrame #

修改__getitem__方法，用iloc按行的位置（而非索引值）取值：

class EmbeddingDataset(torch.utils.data.Dataset):
  def __init__(self,mut_pt, wt_pt, data_df):
    self.pt_mut = mut_pt
    self.pt_wt = wt_pt
    self.df = data_df
  
  def __len__(self):
      return self.pt_mut.shape[0]

  def __getitem__(self, index):
    o1=self.pt_mut[index,:]
    o2=self.pt_wt[index,:]
    if "ddg" in self.df:
      df_out=torch.Tensor([self.df["ddg"].iloc[index]])  # 用iloc按位置取
    else:
      df_out=torch.Tensor([self.df["ID"].iloc[index]])   # 用iloc按位置取
    return  self.pt_mut[index,:],self.pt_wt[index,:],df_out 

额外注意：Dataset参数顺序问题 #

你创建Dataset时的参数顺序和__init__定义不匹配：

# 原代码：__init__参数是(mut_pt, wt_pt, data_df)，但你传的是(train_wt_emb, train_mut_emb, train_df)
train_dataset = EmbeddingDataset(train_wt_emb, train_mut_emb, train_df)

这会导致self.pt_mut实际是wt的embedding，self.pt_wt是mut的embedding，虽然不是当前报错的原因，但会影响后续模型输入，建议修正：

train_dataset = EmbeddingDataset(train_mut_emb, train_wt_emb, train_df)
valid_dataset = EmbeddingDataset(valid_mut_emb, valid_wt_emb, valid_df)

内容的提问来源于stack exchange，提问作者user17230397