29 lines
1.1 KiB
Python
29 lines
1.1 KiB
Python
import torch
|
|
from torch import nn
|
|
from models.convnext_orig import ConvNeXt, CNBlock, CNBlockConfig
|
|
from typing import Any, Callable, List, Optional, Sequence
|
|
|
|
|
|
class DirectModel(nn.Module):
|
|
|
|
def __init__(self, block_setting: List[CNBlockConfig], kernel_size: int = 7, input_features: int = 16, device: str = 'cuda'):
|
|
super(DirectModel, self).__init__()
|
|
|
|
self.encoder = ConvNeXt(block_setting=block_setting, kernel_size=kernel_size, input_features=input_features, device=device)
|
|
self.decoder = ConvNeXt(block_setting=block_setting, kernel_size=kernel_size, device=device, reversed=True)
|
|
|
|
layers = [CNBlock(block_setting[-1].input_channels, kernel_size=kernel_size) for _ in range(5)]
|
|
self.mean_branch = nn.Sequential(*layers).to(device)
|
|
layers = [CNBlock(block_setting[-1].input_channels, kernel_size=kernel_size) for _ in range(5)]
|
|
self.var_branch = nn.Sequential(*layers).to(device)
|
|
|
|
self.flatten_mean = nn.Flatten()
|
|
self.flatten_var = nn.Flatten()
|
|
|
|
|
|
def forward(self, inputs):
|
|
body, template = inputs
|
|
|
|
x = self.encoder(x)
|
|
|
|
return mean |