使用Weights & Biases优化协同过滤算法：超参数调优与实验管理实践

pip install wandb surprise pandas numpy scikit-learn

import wandb
from surprise import Dataset, Reader, SVD
from surprise.model_selection import cross_validate
import pandas as pd

# 初始化W&B项目
wandb.init(project="collaborative-filtering-optimization", 
           config={
               "learning_rate": 0.005,
               "regularization": 0.02,
               "n_factors": 100,
               "n_epochs": 30
           })

# 加载数据
data = pd.read_csv('ratings.csv')
reader = Reader(rating_scale=(1, 5))
surprise_data = Dataset.load_from_df(data[['user_id', 'item_id', 'rating']], reader)

# 基准模型配置
config = wandb.config

# 创建并评估模型
model = SVD(
    lr_all=config.learning_rate,
    reg_all=config.regularization,
    n_factors=config.n_factors,
    n_epochs=config.n_epochs,
    verbose=False
)

# 交叉验证
results = cross_validate(model, surprise_data, measures=['RMSE', 'MAE'], cv=5, verbose=False)

# 记录结果到W&B
wandb.log({
    "rmse": results['test_rmse'].mean(),
    "mae": results['test_mae'].mean(),
    "fit_time": results['fit_time'].mean(),
    "test_time": results['test_time'].mean()
})

wandb.finish()

def train():
    # 初始化W&B运行
    run = wandb.init()
    config = run.config
    
    # 创建模型
    model = SVD(
        lr_all=config.lr,
        reg_all=config.reg,
        n_factors=config.n_factors,
        n_epochs=config.n_epochs,
        verbose=False
    )
    
    # 交叉验证
    results = cross_validate(model, surprise_data, measures=['RMSE'], cv=3, verbose=False)
    rmse = results['test_rmse'].mean()
    
    # 记录结果
    run.log({"rmse": rmse})
    run.finish()
    
    return rmse

# 定义超参数搜索空间
sweep_config = {
    'method': 'bayes',  # 使用贝叶斯优化
    'metric': {
        'name': 'rmse',
        'goal': 'minimize'   
    },
    'parameters': {
        'lr': {
            'distribution': 'log_uniform',
            'min': 1e-5,
            'max': 1e-1
        },
        'reg': {
            'distribution': 'log_uniform', 
            'min': 1e-5,
            'max': 1e-1
        },
        'n_factors': {
            'distribution': 'int_uniform',
            'min': 10,
            'max': 200
        },
        'n_epochs': {
            'distribution': 'int_uniform',
            'min': 10,
            'max': 100
        }
    },
    'early_terminate': {
        'type': 'hyperband',
        'min_iter': 3,
        'max_iter': 10
    }
}

# 创建并启动超参数搜索
sweep_id = wandb.sweep(sweep_config, project="collaborative-filtering-optimization")
wandb.agent(sweep_id, function=train, count=100)

# 连接API获取实验结果
api = wandb.Api()
runs = api.runs("username/collaborative-filtering-optimization")

# 分析最佳运行结果
best_run = min(runs, key=lambda run: run.summary.get('rmse', float('inf')))
print(f"最佳RMSE: {best_run.summary['rmse']}")
print(f"最佳参数: {best_run.config}")

# 可视化参数重要性
import matplotlib.pyplot as plt

# 提取所有运行数据
runs_data = []
for run in runs:
    if 'rmse' in run.summary:
        runs_data.append({
            'rmse': run.summary['rmse'],
            'lr': run.config.get('lr', None),
            'reg': run.config.get('reg', None),
            'n_factors': run.config.get('n_factors', None),
            'n_epochs': run.config.get('n_epochs', None)
        })

# 创建参数相关性分析
df = pd.DataFrame(runs_data)
correlation_matrix = df.corr()

# 记录分析结果
wandb.log({"parameter_correlations": wandb.Table(dataframe=correlation_matrix.reset_index())})

# 保存最佳模型
best_config = best_run.config
best_model = SVD(
    lr_all=best_config['lr'],
    reg_all=best_config['reg'], 
    n_factors=best_config['n_factors'],
    n_epochs=best_config['n_epochs']
)

# 全量训练
trainset = surprise_data.build_full_trainset()
best_model.fit(trainset)

# 创建W&B Artifact保存模型
artifact = wandb.Artifact('best_cf_model', type='model')
with artifact.new_file('model.pkl', mode='wb') as f:
    import pickle
    pickle.dump(best_model, f)

# 添加元数据
artifact.metadata = {
    'algorithm': 'SVD',
    'dataset': 'user_ratings',
    'performance': best_run.summary['rmse']
}

# 记录Artifact
run = wandb.init(project="model-registry")
run.log_artifact(artifact)
run.finish()

# wandb-sweep.yaml
program: train.py
method: bayes
metric:
  name: rmse
  goal: minimize
early_terminate:
  type: hyperband
  min_iter: 3
  max_iter: 10
parameters:
  lr:
    distribution: log_uniform
    min: 0.00001
    max: 0.1
  reg:
    distribution: log_uniform
    min: 0.00001
    max: 0.1
  n_factors:
    distribution: int_uniform
    min: 10
    max: 200
  n_epochs:
    distribution: int_uniform
    min: 10
    max: 100