Spaces:
Sleeping
Sleeping
File size: 3,996 Bytes
2d39721 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 |
"""This module contains the definition of the neural network model used for predictions."""
import torch
class ModuleLayer(torch.nn.Module):
"""Class for the individual layer blocks."""
def __init__(self, intermediate_dim=32, dropout_rate=0.1):
"""Initializer for the 'ModuleLayer' class.
Args:
intermediate_dim (int): The dimension of the intermediate layer.
dropout_rate (float): The dropout rate to apply after the ReLU activation.
"""
super().__init__()
self.mod_linear = torch.nn.Linear(intermediate_dim, intermediate_dim)
self.mod_norm = torch.nn.LayerNorm(normalized_shape=intermediate_dim)
self.mod_relu = torch.nn.ReLU()
self.dropout = torch.nn.Dropout(p=dropout_rate)
def forward(self, x):
"""Forward pass of the layer block.
Args:
x (torch.Tensor): Input tensor.
Returns:
torch.Tensor: Output tensor passing the input through the layer operations.
"""
residual = x
x = self.mod_linear(x)
x = self.mod_norm(x)
x = self.mod_relu(x)
x = self.dropout(x)
x += residual
return x
class Agent(torch.nn.Module):
"""Class for Agent Structure using multiple Layer Blocks."""
def __init__(self, cfg):
"""Initializer for the 'Agent' class.
Args:
cfg (dict): Configuration dictionary containing model parameters.
"""
super().__init__()
self.linear = torch.nn.Linear(
in_features=cfg["in_dim"], out_features=cfg["intermediate_dim"]
)
self.layers = torch.nn.Sequential(
*[
ModuleLayer(
intermediate_dim=cfg["intermediate_dim"], dropout_rate=cfg["dropout_rate"]
)
for _ in range(int(cfg["num_blocks"]))
]
)
self.out = torch.nn.Linear(in_features=cfg["intermediate_dim"], out_features=cfg["out_dim"])
def forward(self, x):
"""Forward pass through the Agent's Layers.
Args:
x (torch.Tensor): Input tensor.
Returns:
x (torch.Tensor): Output tensor after passing through the network.
"""
x = self.linear(x)
x = self.layers(x)
x = self.out(x)
return x
def get_prediction(self, features):
"""Get the deterministic prediction on a single observation or a batch of observations.
Args:
features (torch.tensor): the agent's input features. Expected shape is either
`(num_features,)` for a single observation
or `(batch_size, num_features)` for a batch of observations.
Returns:
prediction_probs (torch.tensor):
- If `features` is a single features (i.e., `features.dim() == 1`), returns a
1-d tensor of the model's probabilities for each decision.
- If `features` is a batch of features (i.e., `features.dim() > 1`),
returns a 2-d tensor, of the model's probabilities for each decision, for the corresponding observation in the batch
"""
# Ensure single samples have a batch dimension
if features.dim() == 1:
# Add a batch dimension if it's a single batch of features
features = features.unsqueeze(0)
with torch.no_grad():
if not isinstance(features, torch.Tensor): # Check if features is not already a tensor
features = torch.tensor(features, dtype=torch.float)
prediction = self.forward(features) # Run a forward pass through the model
if features.size(0) == 1: # This method checks if there is only 1 element in a 1D tensor
return prediction.squeeze() # Returns a single prediction
return prediction # Returns a tensor of predictions
|