import torch
import torch.nn as nn
# 定义因子分解机(Factorization Machine)模型类
class FM(nn.Module):
def __init__(self, dim, k):
super(FM, self).__init__()
self.dim = dim # 特征维度
self.k = k # 因子个数
self.w = nn.Linear(self.dim, 1, bias=True) # 线性部分的权重矩阵
self.v = nn.Parameter(torch.rand(self.dim, self.k) / 100) # 因子分解部分的权重矩阵,通过随机初始化获得
def forward(self, x):
linear = self.w(x) # 计算线性部分的输出
quadratic = 0.5 * torch.sum(torch.pow(torch.mm(x, self.v), 2) - torch.mm(torch.pow(x, 2), torch.pow(self.v, 2))) # 计算因子分解部分的输出
return torch.sigmoid(linear + quadratic) # 返回最终的预测结果
# 训练代码
def train(model, train_loader, optimizer, criterion, epochs):
model.train() # 设置模型为训练模式
for epoch in range(epochs):
running_loss = 0.0
for inputs, labels in train_loader:
optimizer.zero_grad() # 清零梯度
outputs = model(inputs) # 前向传播计算输出
loss = criterion(outputs, labels) # 计算损失函数
loss.backward() # 反向传播计算梯度
optimizer.step() # 更新模型参数
running_loss += loss.item()
epoch_loss = running_loss / len(train_loader)
print(f"Epoch [{epoch + 1}/{epochs}], Loss: {epoch_loss:.4f}")
# 测试代码
def test(model, test_loader):
model.eval() # 设置模型为评估模式
correct = 0
total = 0
with torch.no_grad():
for inputs, labels in test_loader:
outputs = model(inputs) # 前向传播计算输出
predicted = torch.round(outputs) # 对输出进行四舍五入,转为二分类结果(0或1)
total += labels.size(0) # 计算样本总数
correct += (predicted == labels).sum().item() # 统计预测正确的样本数
accuracy = correct / total # 计算准确率
print(f"Accuracy on test set: {accuracy:.4f}")
# 示例数据生成和使用
def main():
# 生成示例数据
batch_size = 32 # 批次大小
dim = 10 # 特征维度
k = 5 # 因子个数
train_data = torch.randn(1000, dim) # 随机生成训练数据
train_labels = torch.randint(0, 2, (1000,)) # 随机生成训练标签(0或1)
test_data = torch.randn(200, dim) # 随机生成测试数据
test_labels = torch.randint(0, 2, (200,)) # 随机生成测试标签(0或1)
train_labels = train_labels.unsqueeze(1) # 将训练标签的维度扩展为(样本数,1)
test_labels = test_labels.unsqueeze(1) # 将测试标签的维度扩展为(样本数,1)
train_labels = train_labels.float() # 将训练标签转换为浮点型
test_labels = test_labels.float() # 将测试标签转换为浮点型
train_dataset = torch.utils.data.TensorDataset(train_data, train_labels) # 构建训练数据集
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True) # 构建训练数据加载器
test_dataset = torch.utils.data.TensorDataset(test_data, test_labels) # 构建测试数据集
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False) # 构建测试数据加载器
# 创建模型实例
model = FM(dim, k)
# 设置优化器、损失函数和训练参数
optimizer = torch.optim.SGD(model.parameters(), lr=0.01) # 使用随机梯度下降(SGD)优化器,学习率为0.01
criterion = nn.BCELoss() # 使用二分类交叉熵作为损失函数
epochs = 100 # 训练轮数
# 训练模型
train(model, train_loader, optimizer, criterion, epochs)
# 在测试集上评估模型
test(model, test_loader)
if __name__ == "__main__":
main()
FM模型的例子
最新推荐文章于 2024-06-13 08:10:28 发布