File size: 1,001 Bytes
5a87d8d |
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 |
import sys
sys.path.append("./BranchSBM")
import torch
import torch.nn as nn
from typing import List, Optional
from networks.mlp_base import SimpleDenseNet
class GeoPathMLP(nn.Module):
def __init__(
self,
input_dim: int,
activation: str,
batch_norm: bool = True,
hidden_dims: Optional[List[int]] = None,
time_geopath: bool = False,
):
super().__init__()
self.input_dim = input_dim
self.time_geopath = time_geopath
self.mainnet = SimpleDenseNet(
input_size=2 * input_dim + (1 if time_geopath else 0),
target_size=input_dim,
activation=activation,
batch_norm=batch_norm,
hidden_dims=hidden_dims,
)
def forward(
self, x0: torch.Tensor, x1: torch.Tensor, t: torch.Tensor
) -> torch.Tensor:
x = torch.cat([x0, x1], dim=1)
if self.time_geopath:
x = torch.cat([x, t], dim=1)
return self.mainnet(x) |