固定输入下PyTorch模型推理成败时GPU内存差异问题

问题现象 #

• 输入数据尺寸固定为[128, 1, 512]，单文件推理时触发CUDA内存不足错误，此时torch.cuda.memory_allocated()为357376、memory_reserved()为2097152
• 对文件夹内所有文件批量推理时完全正常，对应内存值为358400和44040192

单文件推理错误信息 #

...
x.size: torch.Size([128, 1, 512])
x.size: torch.Size([128, 1, 512])
x: tensor([[[0.54997, 0.52163, 0.19563,  ..., 0.34677, 0.23625, 0.21865]],
...
RuntimeError: CUDA out of memory. Tried to allocate 20.00 MiB (GPU 0; 7.78 GiB total capacity; 10.22 MiB already allocated; 1.19 MiB free; 22.00 MiB reserved in total by PyTorch)

批量推理成功输出 #

...
x.size: torch.Size([128, 1, 512])
x.size: torch.Size([128, 1, 512])
x: tensor([[[0.4228, 0.4436, 0.3818,  ..., 0.1387, 0.1879, 0.1545]],
...
Y_pred.shape: (4986,) Y.shape: (4986,)

相关代码 #

批量推理代码 #

model = UC_Model(pModel_filepath="/mnt/confiance/common/UC_Data_Folder/UC_Anomaly_Detection/models/trained_perceiver/best_model_spectra")
(X, Y, X_val, Y_val), (l_data_X, l_data_Y) = model.prepareData(
    pDataPath="/mnt/confiance/common/UC_Data_Folder/UC_Anomaly_Detection/Dataset/test_bench_data/")
X = np.transpose(X, (0, 2, 1))
X_val = np.transpose(X, (0, 2, 1))
print("X.shape: ", X.shape)
print("Y.shape: ", Y.shape)
X_LENGTH = len(X) 
with torch.no_grad():
    correct = 0
    total = 0
    loss = 0
    Y_pred = np.zeros(np.shape(Y))
    for i in range(X_LENGTH // BATCH_SIZE):
        x = torch.from_numpy(
            X[i * BATCH_SIZE:(i + 1) * BATCH_SIZE]
        ).float().to(DEVICE)
        print("x.size: ", x.size())
        print("x.size: ", x.size())
        y = torch.from_numpy(
            Y[i * BATCH_SIZE:(i + 1) * BATCH_SIZE]
        ).long().to(DEVICE)
        x = x.to(DEVICE)
        y_ = model(x)
        y_ = y_.cpu().argmax(dim=-1)
        total += len(y_)
        print("y_.size: ", y_.size(), " y.size(): ", y.size())
        correct += (y_ == y.cpu()).sum().item()
        Y_pred[i * BATCH_SIZE:(i + 1) * BATCH_SIZE] = y_
plot_confusion_matrix(Y, Y_pred, anomaly_list)
print("Y_pred.shape: ", Y_pred.shape, "Y.shape: ", Y.shape)

单文件推理代码 #

def UC_computeInference(pAiModel,pDictParams):
    pNbData=pDictParams['nbData']
    pClasses=pDictParams['classes']
    pDataPathList=pDictParams['dataPathList']
    pDataProd=pDictParams['dataProd']
    pRepertProd=pDictParams['repertProd']
    pSauvegarde=pDictParams['backup']=="oui"
    pListObjets=pDictParams['objets']
    pDataList=pDictParams['dataList']
    listInferError=[]
    x_test=[]
    y_pred=[]
    for i in range(pNbData):
        dataPreprocessed=pAiModel.prepareData(pDataPathList[i])
        prediction=pAiModel(dataPreprocessed)

模型初始化与调用代码 #

class UC_Model:
    #-------------------------------------------------------------------------------
    ## 模型初始化函数
    # @param pModel_filepath : 模型文件路径
    # def __init__(self, pModelpath, pModel_filepath, pParams_filepath="perceiver_params.json"):
    def __init__(self, pModel_filepath):

        self.name="UC_Anomaly_Detection"  # 用例在MODEL_BASE目录下的子目录名称
        self.resPredictionsCibles=[]
        self.resPredictionsSamples=[]
        self.inferenceCible=True

        # 初始化模型
        # self.model = create_perceiver_model()

        # 加载模型
        # /mnt/confiance/common/UC_Data_Folder/UC_Anomaly_Detection/models/trained_perceiver/best_model_spectra 
        self.model = torch.load(pModel_filepath)

    def to_tensor(self, X, Y, i):
        x = torch.from_numpy(
            X[i * BATCH_SIZE:(i + 1) * BATCH_SIZE]
        ).float().to(DEVICE)
        print("x.size: ", x.size()) # 当前x.size: torch.Size([128, 512, 1])，需要改成torch.Size([128, 1, 512])
        x = x.permute(0, 2, 1) # 没关系，我们做维度置换
        print("维度置换后: ")
        print("x.size: ", x.size())
        y = torch.from_numpy(
            Y[i * BATCH_SIZE:(i + 1) * BATCH_SIZE]
        ).long().to(DEVICE)
        x = x.to(DEVICE)
        
        return x, y

    ## 模型调用函数
    # @param pData : 待处理的数据/样本
    # @return 推理结果
    def __call__(self, pData):
        print("pData.size(): ", pData.size())
        print("torch.cuda.memory_allocated(): ", torch.cuda.memory_allocated())
        print("torch.cuda.memory_reserved(): ", torch.cuda.memory_reserved())
        # 用Perceiver模型处理数据
        z = self.model.perceiver(pData)
    
        # 经过线性层
        z = self.model.linear1(z)
        z = z.mean(dim=0)
        z = self.model.linear2(z)
        return F.log_softmax(z, dim=-1)

内容的提问来源于stack exchange，提问作者Revolucion for Monica