BertModel聚类大型访问日志内存溢出，寻求代码优化方案

一、核心问题定位 #

你的代码出现OOM的主要原因有三个：

1. 一次性读取3.5GB日志文件到内存，哪怕只取前1000行，也会先把整个文件加载进来，直接占满内存
2. BERT一次性处理1000条完整日志，生成的张量体积过大，内存峰值飙升
3. 全量相似度矩阵计算和聚类过程的额外内存开销叠加

二、分步优化方案 #

1. 修复大文件读取的内存浪费 #

直接逐行读取并截取前1000行，避免加载整个日志文件：

# 替换原文件读取逻辑
log_lines = []
with open("access.log", "r") as my_file:
    line_count = 0
    while line_count < 1000:
        line = my_file.readline()
        if not line:
            break
        log_lines.append(line.strip())
        line_count += 1

2. 优化BERT嵌入计算的内存占用 #

改用批量处理生成文本嵌入，同时替换为更轻量的DistilBERT模型（参数量比BERT少50%，效果相近）：

from transformers import DistilBertTokenizer, DistilBertModel

# 替换为轻量模型
model_name = 'distilbert-base-uncased'
tokenizer = DistilBertTokenizer.from_pretrained(model_name)
model = DistilBertModel.from_pretrained(model_name)

def get_bert_embeddings(texts, batch_size=32):
    embeddings = []
    model.eval()
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)
    
    with torch.no_grad():
        for i in range(0, len(texts), batch_size):
            batch_texts = texts[i:i+batch_size]
            encoded = tokenizer(
                batch_texts, 
                padding=True, 
                truncation=True, 
                max_length=128,  # 限制文本长度，进一步降低内存
                return_tensors='pt'
            ).to(device)
            
            outputs = model(**encoded)
            batch_embeds = outputs.last_hidden_state[:, 0, :].cpu().numpy()
            embeddings.append(batch_embeds)
            
            # 清理GPU显存（如果使用）
            if device.type == "cuda":
                torch.cuda.empty_cache()
    
    return np.vstack(embeddings)

3. 日志文本预处理（可选但高效） #

日志里的IP、时间戳对聚类无意义，提取核心请求部分（比如GET /path HTTP/1.1），减少BERT处理的文本长度：

import re

def preprocess_log(line):
    # 匹配请求行核心内容
    match = re.search(r'"(GET|POST|PUT|DELETE) [^"]+"', line)
    if match:
        return match.group(0)
    return line

# 读取日志时加入预处理
log_lines = []
with open("access.log", "r") as my_file:
    line_count = 0
    while line_count < 1000:
        line = my_file.readline()
        if not line:
            break
        processed_line = preprocess_log(line.strip())
        if processed_line:
            log_lines.append(processed_line)
            line_count += 1

4. 优化聚类流程 #

用余弦相似度计算替代原有的np.inner，同时简化聚类逻辑：

from sklearn.metrics.pairwise import cosine_similarity

def group_similar_strings(strings):
    similarity_threshold = 0.9
    embeddings = get_bert_embeddings(strings)
    
    # 计算余弦相似度矩阵，转换为距离矩阵
    similarities = cosine_similarity(embeddings)
    distances = 1 - similarities

    clustering = AgglomerativeClustering(
        n_clusters=None, 
        affinity='precomputed', 
        linkage='complete',
        distance_threshold=1 - similarity_threshold
    )
    clusters = clustering.fit_predict(distances)

    groups = defaultdict(list)
    for idx, cluster_id in enumerate(clusters):
        groups[cluster_id].append(strings[idx])
    
    return groups.values()

三、完整优化后代码 #

from transformers import DistilBertTokenizer, DistilBertModel
import numpy as np
import torch
from sklearn.cluster import AgglomerativeClustering
from sklearn.metrics.pairwise import cosine_similarity
from collections import defaultdict
import re

# 初始化轻量模型
model_name = 'distilbert-base-uncased'
tokenizer = DistilBertTokenizer.from_pretrained(model_name)
model = DistilBertModel.from_pretrained(model_name)

def preprocess_log(line):
    # 提取请求核心部分
    match = re.search(r'"(GET|POST|PUT|DELETE) [^"]+"', line)
    if match:
        return match.group(0)
    return line.strip()

def get_bert_embeddings(texts, batch_size=32):
    embeddings = []
    model.eval()
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)
    
    with torch.no_grad():
        for i in range(0, len(texts), batch_size):
            batch_texts = texts[i:i+batch_size]
            encoded = tokenizer(
                batch_texts, 
                padding=True, 
                truncation=True, 
                max_length=128,
                return_tensors='pt'
            ).to(device)
            
            outputs = model(**encoded)
            batch_embeds = outputs.last_hidden_state[:, 0, :].cpu().numpy()
            embeddings.append(batch_embeds)
            
            if device.type == "cuda":
                torch.cuda.empty_cache()
    
    return np.vstack(embeddings)

def group_similar_strings(strings):
    similarity_threshold = 0.9
    embeddings = get_bert_embeddings(strings)
    
    similarities = cosine_similarity(embeddings)
    distances = 1 - similarities

    clustering = AgglomerativeClustering(
        n_clusters=None, 
        affinity='precomputed', 
        linkage='complete',
        distance_threshold=1 - similarity_threshold
    )
    clusters = clustering.fit_predict(distances)

    groups = defaultdict(list)
    for idx, cluster_id in enumerate(clusters):
        groups[cluster_id].append(strings[idx])
    
    return groups.values()

# 读取并预处理前1000行日志
log_lines = []
with open("access.log", "r") as my_file:
    line_count = 0
    while line_count < 1000:
        line = my_file.readline()
        if not line:
            break
        processed_line = preprocess_log(line)
        if processed_line:
            log_lines.append(processed_line)
            line_count += 1

# 执行聚类并打印结果
groups = group_similar_strings(log_lines)
for group in groups:
    print(f"聚类大小: {len(group)}")
    print("样本条目:")
    for entry in group[:3]:
        print(f"  - {entry}")
    print("-"*50)

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