# -*- coding: utf-8 -*-
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: Apache-2.0
# MLP for prediction on the output of readout.
# pylint: disable= no-member, arguments-differ, invalid-name
import torch.nn as nn
# pylint: disable=W0221
"""Two-layer MLP for regression or soft classification
over multiple tasks from graph representations.
For classification tasks, the output will be logits, i.e.
values before sigmoid or softmax.
in_feats : int
Number of input graph features
hidden_feats : int
Number of graph features in hidden layers
n_tasks : int
Number of tasks, which is also the output size.
dropout : float
The probability for dropout. Default to be 0., i.e. no
dropout is performed.
def __init__(self, in_feats, hidden_feats, n_tasks, dropout=0.):
self.predict = nn.Sequential(
[docs] def forward(self, feats):
feats : FloatTensor of shape (B, M3)
* B is the number of graphs in a batch
* M3 is the input graph feature size, must match in_feats in initialization
FloatTensor of shape (B, n_tasks)