models.py

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
import torch.nn.functional as F

class My_MLP(nn.Module):
    def __init__(self, n_layers, in_dim, hidden, out_dim, dropout):
        super(My_MLP, self).__init__()
        self.lins = nn.ModuleList()
        self.lins.append(nn.Linear(in_dim, hidden))
        for _ in range(n_layers - 2):
            self.lins.append(nn.Linear(hidden, hidden))
        self.lins.append(nn.Linear(hidden, out_dim))
        self.dropout = dropout

        # for i in range(len(self.lins)):
        #     nn.init.xavier_uniform_(self.lins[i].weight, gain=nn.init.calculate_gain('relu'))

    def forward(self, features):
        for i, lin in enumerate(self.lins[:-1]):
            features = lin(features)
            features = F.relu(features)
            features = F.dropout(features, p=self.dropout, training=self.training)
        features = self.lins[-1](features)
        return features

class Policy_Net(nn.Module):
    def __init__(self, n_layers, in_dim, hidden, out_dim, dropout):
        super(Policy_Net, self).__init__()
        self.lins = nn.ModuleList()
        self.lins.append(nn.Linear(in_dim, hidden))
        for _ in range(n_layers - 2):
            self.lins.append(nn.Linear(hidden, hidden))
        self.lins.append(nn.Linear(hidden, out_dim))
        self.dropout = dropout

        # for i in range(len(self.lins)):
        #     nn.init.xavier_uniform_(self.lins[i].weight, gain=nn.init.calculate_gain('relu'))

    def forward(self, user, word_embs):
        user = user.unsqueeze(dim=0)
        sentence = word_embs.mean(dim=0, keepdim=True)
        user_sentence = torch.cat([user, sentence], dim=1)
        user_sentence = user_sentence.expand(word_embs.shape[0], -1)
        features = torch.cat([word_embs, user_sentence], dim=1)
        for i, lin in enumerate(self.lins[:-1]):
            features = lin(features)
            features = F.relu(features)
            features = F.dropout(features, p=self.dropout, training=self.training)
        features = self.lins[-1](features)
        return F.softmax(features, dim=0)