class Net(nn.Module):
def __init__(self, args, dev_id):
super(Net, self).__init__()
self._act = get_activation(args.model_activation)
self.encoder = GCMCLayer(args.rating_vals,
args.src_in_units,
args.dst_in_units,
args.gcn_agg_units,
args.gcn_out_units,
args.gcn_dropout,
args.gcn_agg_accum,
agg_act=self._act,
share_user_item_param=args.share_param,
device=dev_id)
if args.mix_cpu_gpu and args.use_one_hot_fea:
# if use_one_hot_fea, user and movie feature is None
# W can be extremely large, with mix_cpu_gpu W should be stored in CPU
self.encoder.partial_to(dev_id)
else:
self.encoder.to(dev_id)
self.decoder = BiDecoder(in_units=args.gcn_out_units,
num_classes=len(args.rating_vals),
num_basis=args.gen_r_num_basis_func)
self.decoder.to(dev_id)
def forward(self, compact_g, frontier, ufeat, ifeat,
possible_rating_values):
user_out, movie_out = self.encoder(frontier, ufeat, ifeat)
pred_ratings = self.decoder(compact_g, user_out, movie_out)
return pred_ratings
class Net(nn.Module):
def __init__(self, args, dev_id):
super(Net, self).__init__()
self._act = get_activation(args.model_activation)
self.encoder = GCMCLayer(args.rating_vals,
args.src_in_units,
args.dst_in_units,
args.gcn_agg_units,
args.gcn_out_units,
args.gcn_dropout,
args.gcn_agg_accum,
agg_act=self._act,
share_user_item_param=args.share_param,
device=dev_id)
if args.mix_cpu_gpu and args.use_one_hot_fea:
# if use_one_hot_fea, user and movie feature is None
# W can be extremely large, with mix_cpu_gpu W should be stored in CPU
self.encoder.partial_to(dev_id)
else:
self.encoder.to(dev_id)
self.decoder = DenseBiDecoder(in_units=args.gcn_out_units,
num_classes=len(args.rating_vals),
num_basis=args.gen_r_num_basis_func)
self.decoder.to(dev_id)
def forward(self, compact_g, frontier, ufeat, ifeat,
possible_rating_values):
user_out, movie_out = self.encoder(frontier, ufeat, ifeat)
head_emb = []
tail_emb = []
for possible_rating_value in possible_rating_values:
head, tail = compact_g.all_edges(etype=str(possible_rating_value))
head_emb.append(user_out[head])
tail_emb.append(movie_out[tail])
head_emb = th.cat(head_emb, dim=0)
tail_emb = th.cat(tail_emb, dim=0)
pred_ratings = self.decoder(head_emb, tail_emb)
return pred_ratings
