Spaces:
Build error
Build error
import torch | |
import torch.nn as nn | |
from fastai.vision import * | |
from modules.model import _default_tfmer_cfg | |
from modules.resnet import resnet45 | |
from modules.transformer import (PositionalEncoding, | |
TransformerEncoder, | |
TransformerEncoderLayer) | |
class ResTranformer(nn.Module): | |
def __init__(self, config): | |
super().__init__() | |
alpha_d = ifnone(config.model_vision_backbone_alpha_d, 1.) | |
self.d_model = ifnone(config.model_vision_d_model, _default_tfmer_cfg['d_model']) | |
self.resnet = resnet45(alpha_d, output_channels=self.d_model) | |
nhead = ifnone(config.model_vision_nhead, _default_tfmer_cfg['nhead']) | |
d_inner = ifnone(config.model_vision_d_inner, _default_tfmer_cfg['d_inner']) | |
dropout = ifnone(config.model_vision_dropout, _default_tfmer_cfg['dropout']) | |
activation = ifnone(config.model_vision_activation, _default_tfmer_cfg['activation']) | |
num_layers = ifnone(config.model_vision_backbone_ln, 2) | |
self.pos_encoder = PositionalEncoding(self.d_model, max_len=8*32) | |
encoder_layer = TransformerEncoderLayer(d_model=self.d_model, nhead=nhead, | |
dim_feedforward=d_inner, dropout=dropout, activation=activation) | |
self.transformer = TransformerEncoder(encoder_layer, num_layers) | |
def forward_transformer(self, feature): | |
n, c, h, w = feature.shape | |
feature = feature.view(n, c, -1).permute(2, 0, 1) | |
feature = self.pos_encoder(feature) | |
feature = self.transformer(feature) | |
feature = feature.permute(1, 2, 0).view(n, c, h, w) | |
return feature | |
def forward(self, images, **kwargs): | |
feature = self.resnet(images, **kwargs) | |
feature = self.forward_transformer(feature) | |
return feature |