Google Colab RAM不足崩溃，排查torch.tensor().squeeze(0)异常

问题根源：编码环节出错 #

你的RAM溢出问题核心出在编码函数的错误处理，而非张量构造环节：

1. 错误地将row转为list(row)，这会把句子拆成单个字符，让tokenizer对每个字符单独编码，导致每个样本生成的token数量暴增（比如100字符的句子会生成100组token序列）。
2. tokenizer.encode方法不支持axis=1参数，这个无效参数会导致编码结果的结构异常，后续转张量时生成远超预期的超大维度数据，直接占满RAM。
3. 未显式设置max_length，默认会使用模型的最大输入长度（RoBERTa为512），进一步加剧内存消耗。

修复代码（无需Colab Pro） #

方案1：修正单个样本编码逻辑 #

model = RobertaForSequenceClassification.from_pretrained(model_name)
tokenizer = RobertaTokenizer.from_pretrained(model_name)

def encode_df(dataframe):
    # 直接传入字符串，不要拆成字符列表；显式设置max_length控制token数量
    dataframe['sentences'] = dataframe['sentences'].apply(
        lambda row: tokenizer.encode(
            row,
            add_special_tokens=True,
            padding='max_length',
            truncation=True,
            max_length=128  # 根据任务调整，比如256，不要超过模型最大限制
        )
    )
    return dataframe

df = pd.read_excel("/content/profile sentences smaller set.xlsx")
profiles_encoded = encode_df(df)
profiles_input = profiles_encoded['sentences'].values.tolist()

input_ids = torch.tensor(profiles_input)
# 仅当你确认张量有多余的维度时再用squeeze(0)，否则不要随意压缩维度

方案2：用批量编码提升效率（更推荐） #

批量编码能减少中间内存占用，还能直接返回tensor，省去手动转换步骤：

model = RobertaForSequenceClassification.from_pretrained(model_name)
tokenizer = RobertaTokenizer.from_pretrained(model_name)

df = pd.read_excel("/content/profile sentences smaller set.xlsx")

# 直接批量编码，返回tensor
encoded_data = tokenizer.batch_encode_plus(
    df['sentences'].tolist(),
    add_special_tokens=True,
    padding='max_length',
    truncation=True,
    max_length=128,
    return_tensors='pt'
)
input_ids = encoded_data['input_ids']

额外内存优化技巧 #

• 释放GPU缓存：在代码中加入torch.cuda.empty_cache()，定期清理无用的张量内存。
• 换用轻量模型：如果roberta-base内存压力大，换成distilroberta-base（参数减少60%，性能损失小）或roberta-small。
• 避免冗余操作：不要重复创建DataFrame（你的代码里df = pd.DataFrame(df)是多余的），减少内存占用。

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