|
|
import torch
|
|
|
import numpy as np
|
|
|
from AutoRec.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._model = model
|
|
|
self._config = config
|
|
|
self._optimizer = pick_optimizer(self._model, self._config)
|
|
|
|
|
|
def _train_single_batch(self, batch_x, batch_mask_x):
|
|
|
"""
|
|
|
对单个小批量数据进行训练
|
|
|
"""
|
|
|
if self._config['use_cuda'] is True:
|
|
|
# 将这些数据由CPU迁移到GPU
|
|
|
batch_x, batch_mask_x = batch_x.cuda(), batch_mask_x.cuda()
|
|
|
|
|
|
# 模型的输入为用户评分向量或者物品评分向量，调用forward进行前向传播
|
|
|
ratings_pred = self._model(batch_x.float())
|
|
|
# 通过交叉熵损失函数来计算损失, ratings_pred.view(-1)代表将预测结果摊平，变成一维的结构。
|
|
|
loss, rmse = self._model.loss(res=ratings_pred, input=batch_x, mask=batch_mask_x, optimizer=self._optimizer)
|
|
|
# 先将梯度清零,如果不清零，那么这个梯度就和上一个mini-batch有关
|
|
|
self._optimizer.zero_grad()
|
|
|
# 反向传播计算梯度
|
|
|
loss.backward()
|
|
|
# 梯度下降等优化器 更新参数
|
|
|
self._optimizer.step()
|
|
|
# 将loss的值提取成python的float类型
|
|
|
loss = loss.item()
|
|
|
return loss, rmse
|
|
|
|
|
|
def _train_an_epoch(self, train_loader, epoch_id, train_mask):
|
|
|
"""
|
|
|
训练一个Epoch，即将训练集中的所有样本全部都过一遍
|
|
|
"""
|
|
|
# 告诉模型目前处于训练模式，启用dropout以及batch normalization
|
|
|
self._model.train()
|
|
|
total_loss = 0
|
|
|
total_rmse = 0
|
|
|
# 从DataLoader中获取小批量的id以及数据
|
|
|
for batch_id, (batch_x, batch_mask_x) in enumerate(train_loader):
|
|
|
assert isinstance(batch_x, torch.Tensor)
|
|
|
assert isinstance(batch_mask_x, torch.Tensor)
|
|
|
|
|
|
loss, rmse = self._train_single_batch(batch_x, batch_mask_x)
|
|
|
print('[Training Epoch: {}] Batch: {}, Loss: {}, RMSE: {}'.format(epoch_id, batch_id, loss, rmse))
|
|
|
total_loss += loss
|
|
|
total_rmse += rmse
|
|
|
rmse = np.sqrt(total_rmse.detach().cpu().numpy() / (train_mask == 1).sum())
|
|
|
print('Training Epoch: {}, Total Loss: {}, total RMSE: {}'.format(epoch_id, total_loss, rmse))
|
|
|
|
|
|
def train(self, train_r, train_mask_r):
|
|
|
# 是否使用GPU加速
|
|
|
self.use_cuda()
|
|
|
|
|
|
for epoch in range(self._config['num_epoch']):
|
|
|
print('-' * 20 + ' Epoch {} starts '.format(epoch) + '-' * 20)
|
|
|
# 构造一个DataLoader
|
|
|
data_loader = Construct_DataLoader(train_r, train_mask_r, batchsize=self._config['batch_size'])
|
|
|
# 训练一个轮次
|
|
|
self._train_an_epoch(data_loader, epoch_id=epoch, train_mask=train_mask_r)
|
|
|
|
|
|
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 save(self):
|
|
|
self._model.saveModel()
|
