44 lines
1.2 KiB
Python
44 lines
1.2 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.optim as optim
|
|
import pandas as pd
|
|
from sklearn.preprocessing import StandardScaler
|
|
|
|
class RecommendModel(nn.Module):
|
|
def __init__(self, input_size, embedding_size):
|
|
super().__init__()
|
|
self.encoder = nn.Sequential(
|
|
nn.Linear(input_size, 256),
|
|
nn.ReLU(),
|
|
nn.Dropout(0.3),
|
|
nn.Linear(256, 128),
|
|
nn.ReLU(),
|
|
nn.Linear(128, embedding_size)
|
|
)
|
|
|
|
def forward(self, x):
|
|
return self.encoder(x)
|
|
|
|
def train_model():
|
|
# 加载数据
|
|
videos_df = pd.read_csv('videos.csv')
|
|
behaviors_df = pd.read_csv('behaviors.csv')
|
|
|
|
# 数据预处理
|
|
scaler = StandardScaler()
|
|
features = scaler.fit_transform(videos_df[['duration', 'view_count', 'like_count']])
|
|
|
|
# 创建模型
|
|
model = RecommendModel(input_size=features.shape[1], embedding_size=128)
|
|
criterion = nn.CosineEmbeddingLoss()
|
|
optimizer = optim.Adam(model.parameters())
|
|
|
|
# 训练模型
|
|
for epoch in range(100):
|
|
# ... 训练代码
|
|
|
|
# 保存模型
|
|
torch.save(model.state_dict(), '../src/main/resources/models/recommend_model.pt')
|
|
|
|
if __name__ == '__main__':
|
|
train_model() |