File size: 966 Bytes
3ef1661 |
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 |
import torch
import torch.nn as nn
from mono.utils.comm import get_func
class DensePredModel(nn.Module):
def __init__(self, cfg):
super(DensePredModel, self).__init__()
self.encoder = get_func('mono.model.' + cfg.model.backbone.prefix + cfg.model.backbone.type)(**cfg.model.backbone)
self.decoder = get_func('mono.model.' + cfg.model.decode_head.prefix + cfg.model.decode_head.type)(cfg)
# try:
# decoder_compiled = torch.compile(decoder, mode='max-autotune')
# "Decoder compile finished"
# self.decoder = decoder_compiled
# except:
# "Decoder compile failed, use default setting"
# self.decoder = decoder
self.training = True
def forward(self, input, **kwargs):
# [f_32, f_16, f_8, f_4]
features = self.encoder(input)
# [x_32, x_16, x_8, x_4, x, ...]
out = self.decoder(features, **kwargs)
return out |