CLYang617/RemoteSensingChangeDetection-RSCD.HA2F

1import torch2import torch.nn as nn3 4from model.encoder import Encoder5from model.decoder import Decoder6 7from model.utils import weight_init8 9 10class Trainer(nn.Module):11    def __init__(self, model_type='small'):12        super().__init__()13        if model_type == 'tiny':14            embed_dim = 19215        elif model_type == 'small':16            embed_dim = 38417        else:18            assert False, r'Trainer: check the vit model type'19 20        self.encoder = Encoder(model_type)21 22        self.decoder = Decoder(in_dim=[64, 128, 256, embed_dim])23        weight_init(self.decoder)24        25    def forward(self, x, y):26        fx, fy = self.encoder(x, y)27        pred = self.decoder(fx, fy)28 29        return pred30