[D2L] DeepFM(code review)

class DeepFM(nn.Module):
    '''
    input: features
    output : 0~1사이의 실수
    '''
    def __init__(self, field_dims, num_factors, mlp_dims, drop_rate=0.1):
        super(DeepFM, self).__init__()
        self.num_inputs = int(sum(field_dims))
        self.embedding = nn.Embedding(self.num_inputs, num_factors)
        self.fc = nn.Embedding(self.num_inputs, 1)
        self.linear_layer = nn.Linear(1, 1)
        input_dim = self.embed_output_dim = len(field_dims) * num_factors
        self.mlp = []
        for dim in mlp_dims:
            self.mlp.append(nn.Linear(input_dim, dim))
            self.mlp.append(nn.ReLU())
            self.mlp.append(nn.Dropout(p=drop_rate))
            input_dim = dim
        self.mlp.append(nn.Linear(input_dim, 1))
        self.mlp = nn.Sequential(*self.mlp)

    def forward(self, x):
        # x : (batch_size, feature개수)
        embed_x = self.embedding(x) # embed_x : (batch_size, feature개수, num_factors)
        square_of_sum = torch.sum(embed_x, axis=1) ** 2 # (batch_size, num_factors)
        sum_of_square = torch.sum(embed_x ** 2, axis=1) # (batch_size, num_factors)
        inputs = embed_x.view(-1, self.embed_output_dim) # (batch_size, feature개수*num_factors)
        # fc(x) : (batch_size, feature개수, 1)
        x = self.linear_layer(self.fc(x).sum(1)) \
            + 0.5 * (square_of_sum - sum_of_square).sum(1, keepdims=True) \
            + self.mlp(inputs)
        x = torch.sigmoid(x)
        return x

embedding layer를 이용해 모든 피쳐를 같은 차원을 가지도록 변환

이를 sum of square, square of sum, mlp 세 파트로 진행

참고로 앞의 두개의 구조에서 여러 피쳐들을 sum하는 것이기때문에 batch size를 고려하지 않는다면 shape은 (1, num_factors) (*num_factors는 임베딩 차원을 의미) 를 가진다.

후에 두개의 term간의 연산을 한다음 sum을 하면 하나의 스칼라값이 나온다.

이를 linear layer의 결과와 mlp단의 결과와 합하면 최종 예측값이 나오게된다.