21 lines
690 B
Python
21 lines
690 B
Python
import torch
|
|
from torch import nn
|
|
from models.base import AEEncoder, Decoder
|
|
|
|
|
|
class AutoEncoder(nn.Module):
|
|
|
|
def __init__(self, config):
|
|
super(AutoEncoder, self).__init__()
|
|
|
|
self.encoder : AEEncoder = AEEncoder(config['encoder_block_settings'], config['kernel_size'], config['activation'], is_label=config['is_label'], device=config['device'], nr_of_labels=7)
|
|
self.decoder : Decoder = Decoder(config['decoder_block_settings'], kernel_size=config['kernel_size'], is_label=config['is_label'], device=config['device'], nr_of_classes=6)
|
|
|
|
|
|
def forward(self, input):
|
|
x = input
|
|
|
|
x = self.encoder(x)
|
|
x = self.decoder(x)
|
|
|
|
return x |