Spaces:
Running on Zero
Running on Zero
| """ The code below is an adaptation of DenseDepth (https://github.com/ialhashim/DenseDepth) | |
| We replaced the BatchNorm with GroupNorm, | |
| added weight standardization, | |
| added convex upsampling layer, | |
| and allowed the backbone and decoder depth to be changeable | |
| """ | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from models.conv_encoder_decoder.submodules import Encoder, Decoder | |
| class DenseDepth(nn.Module): | |
| def __init__(self, num_classes, | |
| B=5, pretrained=True, | |
| NF=2048, BN=True, | |
| down=2, learned_upsampling=True, | |
| **kwargs): | |
| super(DenseDepth, self).__init__() | |
| self.encoder = Encoder(B=B, pretrained=pretrained) | |
| self.decoder = Decoder(num_classes=num_classes, | |
| B=B, NF=NF, BN=BN, | |
| down=down, learned_upsampling=learned_upsampling, | |
| **kwargs) | |
| def forward(self, x, **kwargs): | |
| return self.decoder(self.encoder(x), **kwargs) | |
| def get_1x_lr_params(self): # lr/10 learning rate | |
| return self.encoder.parameters() | |
| def get_10x_lr_params(self): # lr learning rate | |
| modules = [self.decoder] | |
| for m in modules: | |
| yield from m.parameters() | |