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