|
|
import torch
|
|
|
from NCF.dataloader import Construct_DataLoader
|
|
|
|
|
|
def pick_optimizer(network, params):
|
|
|
optimizer = None
|
|
|
if params['optimizer'] == 'sgd':
|
|
|
optimizer = torch.optim.SGD(network.parameters(),
|
|
|
lr=params['sgd_lr'],
|
|
|
momentum=params['sgd_momentum'],
|
|
|
weight_decay=params['l2_regularization'])
|
|
|
elif params['optimizer'] == 'adam':
|
|
|
optimizer = torch.optim.Adam(network.parameters(),
|
|
|
lr=params['adam_lr'],
|
|
|
weight_decay=params['l2_regularization'])
|
|
|
elif params['optimizer'] == 'rmsprop':
|
|
|
optimizer = torch.optim.RMSprop(network.parameters(),
|
|
|
lr=params['rmsprop_lr'],
|
|
|
alpha=params['rmsprop_alpha'],
|
|
|
momentum=params['rmsprop_momentum'])
|
|
|
return optimizer
|
|
|
|
|
|
class Trainer(object):
|
|
|
def __init__(self, model, config):
|
|
|
self._config = config
|
|
|
self._model = model
|
|
|
# 选择优化器
|
|
|
self._optimizer = pick_optimizer(self._model, self._config)
|
|
|
# 定义损失函数，对于隐反馈数据，这里使用交叉熵损失函数
|
|
|
self._crit = torch.nn.BCELoss()
|
|
|
|
|
|
def _train_single_batch(self, users, items, ratings):
|
|
|
"""
|
|
|
对单个小批量数据进行训练
|
|
|
:param users: user Tensor
|
|
|
:param items: item Tensor
|
|
|
:param ratings: rating Tensor
|
|
|
:return:
|
|
|
"""
|
|
|
if self._config['use_cuda'] is True:
|
|
|
# 将这些数据由CPU迁移到GPU
|
|
|
users, items, ratings = users.cuda(), items.cuda(), ratings.cuda()
|
|
|
|
|
|
# 先将梯度清零,如果不清零，那么这个梯度就和上一个mini-batch有关
|
|
|
self._optimizer.zero_grad()
|
|
|
# 模型的输入users， items，调用forward进行前向传播
|
|
|
ratings_pred = self._model(users, items)
|
|
|
# 通过交叉熵损失函数来计算损失, ratings_pred.view(-1)代表将预测结果摊平，变成一维的结构。
|
|
|
loss = self._crit(ratings_pred.view(-1), ratings)
|
|
|
# 反向传播计算梯度
|
|
|
loss.backward()
|
|
|
# 梯度下降等优化器 更新参数
|
|
|
self._optimizer.step()
|
|
|
# 将loss的值提取成python的float类型
|
|
|
loss = loss.item()
|
|
|
return loss
|
|
|
|
|
|
def _train_an_epoch(self, train_loader, epoch_id):
|
|
|
"""
|
|
|
训练一个Epoch，即将训练集中的所有样本全部都过一遍
|
|
|
:param train_loader: Torch的DataLoader
|
|
|
:param epoch_id: 训练轮次Id
|
|
|
:return:
|
|
|
"""
|
|
|
# 告诉模型目前处于训练模式，启用dropout以及batch normalization
|
|
|
self._model.train()
|
|
|
total_loss = 0
|
|
|
# 从DataLoader中获取小批量的id以及数据
|
|
|
for batch_id, batch in enumerate(train_loader):
|
|
|
assert isinstance(batch[0], torch.LongTensor)
|
|
|
# 这里的user, item, rating大小变成了1024维了，因为batch_size是1024，即每次选取1024个样本数据进行训练
|
|
|
user, item, rating = batch[0], batch[1], batch[2]
|
|
|
rating = rating.float()
|
|
|
loss = self._train_single_batch(user, item, rating)
|
|
|
print('[Training Epoch {}] Batch {}, Loss {}'.format(epoch_id, batch_id, loss))
|
|
|
total_loss += loss
|
|
|
print('Training Epoch: {}, TotalLoss: {}'.format(epoch_id, total_loss))
|
|
|
|
|
|
def train(self, sampleGenerator):
|
|
|
# 是否使用GPU加速
|
|
|
self.use_cuda()
|
|
|
# 是否使用预先训练好的参数
|
|
|
self.load_preTrained_weights()
|
|
|
|
|
|
for epoch in range(self._config['num_epoch']):
|
|
|
print('-' * 20 + ' Epoch {} starts '.format(epoch) + '-' * 20)
|
|
|
# 每个轮次都重新随机产生样本数据集
|
|
|
users, items, ratings = sampleGenerator(num_negatives=self._config['num_negative'])
|
|
|
# 构造一个DataLoader
|
|
|
data_loader = Construct_DataLoader(users=users, items=items, ratings=ratings,
|
|
|
batchsize=self._config['batch_size'])
|
|
|
# 训练一个轮次
|
|
|
self._train_an_epoch(data_loader, epoch_id=epoch)
|
|
|
|
|
|
def use_cuda(self):
|
|
|
if self._config['use_cuda'] is True:
|
|
|
assert torch.cuda.is_available(), 'CUDA is not available'
|
|
|
torch.cuda.set_device(self._config['device_id'])
|
|
|
self._model.cuda()
|
|
|
|
|
|
def load_preTrained_weights(self):
|
|
|
if self._config['pretrain'] is True:
|
|
|
self._model.load_preTrained_weights()
|
|
|
|
|
|
def save(self):
|
|
|
self._model.saveModel()
|
