PyTorch前向损失函数内拆分增强后Tensor/Batch遇问题求助

你的思路完全正确——就是要把增强后的512个样本拆分为前256张（原始未增强图）和后256张（对应增强图），再分别计算特征并对比分数。

问题出在拆分维度选错了：你之前是在特征维度（第二个维度，dim=1）拆分[512,128]的张量，所以得到两个[512,64]的结果；但正确的做法是在**batch维度（第一个维度，dim=0）**进行拆分。

正确代码示例 #

假设卷积后的特征张量为features，形状是[512, 128]，执行以下代码即可得到预期的两个张量：

import torch

# 按batch维度拆分，分成前256个和后256个样本的特征
original_features, augmented_features = torch.split(features, 256, dim=0)

拆分后original_features和augmented_features的形状都会是[256, 128]，完全符合你的需求。

补充说明 #

从你给出的增强后images张量形状[512,3,96,96]来看，你的数据增强逻辑应该是把每个原始样本生成一个增强样本，然后按[原图1, 原图2, ..., 原图256, 增强图1, 增强图2, ..., 增强图256]的顺序拼接成了512的batch，所以按dim=0拆分前256和后256就能精准对应到原图和其对应的增强图。

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