如何用unsqueeze得到多分类标签的[37715,1]Torch张量？Dataset优化建议

标签形状异常的原因及修复 #

核心问题出在对输入特征X的错误维度扩展：

self.X = torch.tensor(self.X).float().unsqueeze(0)

这一步给整个特征矩阵额外增加了一个维度，导致self.X的形状变成[1, 171812, 384]，len(self.X)返回1，所以__getitem__(0)取出的是整个特征集合，对应的标签y[0]自然只是单个标签（形状[1]）。

修复步骤：

1. 去掉unsqueeze(0)，让self.X保持原始样本维度：

self.X = torch.tensor(stored_embeddings).float()  # 形状应为[37715, 384]

2. 标签部分调整为分类任务适配的整数类型，保留维度扩展：

self.y = torch.tensor(stored_labels.to_numpy()).long().unsqueeze(-1)  # 形状[37715, 1]

3. 修复后len(self.X)会返回37715，调用train_data[idx]时，单个样本X形状为[384]，y形状为[1]；使用DataLoader加载时，batch数据会自动组合成[batch_size, 384]和[batch_size, 1]，符合预期。

EmbeddingDataset类优化建议 #

• 启用transform参数：当前代码定义了transform和target_transform但未使用，需在__getitem__中添加转换逻辑：

  def __getitem__(self, idx):
      X = self.X[idx]
      y = self.y[idx]
      if self.transform:
          X = self.transform(X)
      if self.target_transform:
          y = self.target_transform(y)
      return X, y
  

• 简化冗余赋值：将self.X的赋值合并为一行，减少内存占用：

  self.X = torch.tensor(stored_embeddings, dtype=torch.float32)
  

• 标签类型适配任务：多分类任务标签需用torch.long类型（对应整数类别），而非float，适配PyTorch内置的CrossEntropyLoss等损失函数。
• 增加异常处理：对文件读取、pickle加载过程添加异常捕获，提升鲁棒性：

  def __init__(self, embedding_fp, transform=None, target_transform=None):
      try:
          with open(embedding_fp, "rb") as fIn:
              stored_data = pickle.load(fIn)
              stored_labels = stored_data['labels']
              stored_embeddings = stored_data['embeddings']
      except FileNotFoundError:
          raise ValueError(f"文件 {embedding_fp} 不存在")
      except pickle.UnpicklingError:
          raise ValueError(f"无法解析 {embedding_fp}，文件可能损坏")
      # 后续数据处理逻辑
  

• 校验样本一致性：初始化时检查特征与标签的样本数量是否匹配：

  if len(stored_embeddings) != len(stored_labels):
      raise ValueError("特征与标签的样本数量不匹配")
  

• 延迟数据转换（可选）：若数据集过大，可在__getitem__中才将numpy数组转为tensor，降低初始化阶段的内存占用：

  def __init__(self, embedding_fp, transform=None, target_transform=None):
      with open(embedding_fp, "rb") as fIn:
          stored_data = pickle.load(fIn)
          self.y = stored_labels.to_numpy()
          self.X = stored_embeddings
      # 其他初始化逻辑
  def __getitem__(self, idx):
      X = torch.tensor(self.X[idx], dtype=torch.float32)
      y = torch.tensor(self.y[idx], dtype=torch.long).unsqueeze(-1)
      # 应用transform逻辑
      return X, y
  

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