def set_params(self):
for class_path, params in self.class_params.iteritems():
klass= import_a_thing(class_path)
for k, v in params.iteritems():
if not isinstance(v, basestring): continue
if v.isdigit(): params[k]= int(v)
elif re.match('[0-9]+.[0-9]+', v): params[k]= float(v)
set_params(klass, params)
elif args.output_scale == 32:
self.encoder = resnext101_32x8d(args.use_pretrain)
else:
raise BaseException("output scale should be 16 or 32")
self.aspp = ASPP(args)
self.decoder = Decoder(args)
if args.freeze_bn:
for m in self.modules():
if isinstance(m, nn.BatchNorm2d):
m.eval()
def forward(self, input):
# input: batch_size, channel, height, width
output, low_level_feature = self.encoder(input)
output = self.aspp(output)
output, loss_sigma = self.decoder(output, low_level_feature)
# print(output[0].size())
return output, loss_sigma
if __name__ == '__main__':
from utils.params import set_params
args = set_params()
model = AttDepth(args)
model.eval()
input = torch.rand(1, 3, 320, 480)
output = model(input)
print(output.size())