首先,我们需要导入所需的库:
import torch
import torch.nn as nn
import torch.optim as optim
然后,我们定义一个类来实现最小的推荐算法:
class RecommendationModel(nn.Module):
def __init__(self, num_users, num_items, embedding_dim):
super(RecommendationModel, self).__init__()
self.user_embedding = nn.Embedding(num_users, embedding_dim)
self.item_embedding = nn.Embedding(num_items, embedding_dim)
self.fc = nn.Linear(embedding_dim, 1)
def forward(self, users, items):
user_embedded = self.user_embedding(users)
item_embedded = self.item_embedding(items)
prediction = self.fc(torch.mul(user_embedded, item_embedded)).squeeze()
return prediction
在上述代码中,我们定义了一个继承自nn.Module的类RecommendationModel。在初始化函数中,我们定义了两个嵌入层(Embedding)以及一个全连接层(Linear)。forward函数实现了模型的前向传播过程,其中计算预测值的方法是对用户和物品的嵌入向量进行元素级别的乘法操作,并通过全连接层得到最终的预测值。
接下来,我们可以定义训练函数:
def train_model(model, train_data, num_epochs, learning_rate):
optimizer = optim.Adam(model.parameters(), lr=learning_rate)
criterion = nn.MSELoss()
for epoch in range(num_epochs):
total_loss = 0
for users, items, ratings in train_data:
optimizer.zero_grad()
predictions = model(users, items)
loss = criterion(predictions, ratings)
loss.backward()
optimizer.step()
total_loss += loss.item()
print("Epoch {}/{} Loss: {:.4f}".format(epoch+1, num_epochs, total_loss))
在上述代码中,我们使用Adam优化器和均方误差损失函数(MSELoss)来进行模型的训练。对于每个epoch,我们计算总的损失,并在每个iteration中进行反向传播和参数更新。
最后,我们可以使用上述定义的函数来训练模型:
# 假设有100个用户和200个物品,嵌入维度为10
num_users = 100
num_items = 200
embedding_dim = 10
# 生成随机训练数据
train_data = [(torch.randint(num_users, (1,)), torch.randint(num_items, (1,)), torch.rand(1)) for _ in range(1000)]
# 创建模型实例
model = RecommendationModel(num_users, num_items, embedding_dim)
# 训练模型
num_epochs = 10
learning_rate = 0.001
train_model(model, train_data, num_epochs, learning_rate)
在上述代码中,我们通过torch.randint函数生成1000个随机的用户、物品和评分数据作为训练数据。然后,我们创建了一个模型实例,并使用train_model函数对模型进行训练。