def define_model(encoder='resnet'): if encoder is 'resnet': original_model = resnet.resnet50(pretrained=True) Encoder = modules.E_resnet(original_model) model = net.model(Encoder, num_features=2048, block_channel=[256, 512, 1024, 2048]) if encoder is 'densenet': original_model = densenet.densenet161(pretrained=True) Encoder = modules.E_densenet(original_model) model = net.model(Encoder, num_features=2208, block_channel=[192, 384, 1056, 2208]) if encoder is 'senet': original_model = senet.senet154(pretrained='imagenet') Encoder = modules.E_senet(original_model) model = net.model(Encoder, num_features=2048, block_channel=[256, 512, 1024, 2048]) if encoder is 'resnet4': original_model = resnet4.resnet50(pretrained=True) Encoder = modules.E_resnet(original_model) model = net.model(Encoder, num_features=2048, block_channel=[256, 512, 1024, 2048]) return model
def define_model(is_resnet, is_densenet, is_senet): use18 = True # True if is_resnet: if not use18: original_model = resnet.resnet18(pretrained = True) Encoder = modules.E_resnet(original_model) model = net.model(Encoder, num_features=512, block_channel = [64, 128, 256, 512]) else: stereoModel = Resnet18Encoder(3) model_dict = stereoModel.state_dict() encoder_dict = torch.load('./models/monodepth_resnet18_001.pth',map_location='cpu' ) new_dict = {} for key in encoder_dict: if key in model_dict: new_dict[key] = encoder_dict[key] stereoModel.load_state_dict(new_dict ) Encoder = stereoModel model = net.model(Encoder, num_features=512, block_channel = [64, 128, 256, 512]) if is_densenet: original_model = densenet.densenet161(pretrained=True) Encoder = modules.E_densenet(original_model) model = net.model(Encoder, num_features=2208, block_channel = [192, 384, 1056, 2208]) if is_senet: original_model = senet.senet154(pretrained='imagenet') Encoder = modules.E_senet(original_model) model = net.model(Encoder, num_features=2048, block_channel = [256, 512, 1024, 2048]) return model
def define_model(): original_model = resnet.resnet50(pretrained=True) Encoder = modules.E_resnet(original_model) model = net.model(None, Encoder, num_features=2048, block_channel=[256, 512, 1024, 2048]) return model
def define_model(is_resnet, is_densenet, is_senet): if is_resnet: original_model = resnet.resnet50(pretrained=True) Encoder = modules.E_resnet(original_model) model = net.model(Encoder, num_features=2048, block_channel=[256, 512, 1024, 2048]) if is_densenet: original_model = densenet.densenet161(pretrained=True) Encoder = modules.E_densenet(original_model) model = net.model(Encoder, num_features=2208, block_channel=[192, 384, 1056, 2208]) if is_senet: original_model = senet.senet154(pretrained=None) Encoder = modules.E_senet(original_model) model = net.model(Encoder, num_features=2048, block_channel=[256, 512, 1024, 2048]) return model
def define_model(is_resnet, is_densenet, is_senet, model='tbdp', parallel=False, semff=False, pcamff=False): if is_resnet: original_model = resnet.resnet50(pretrained=True) Encoder = modules.E_resnet(original_model) if model == 'tbdp': model = net.TBDPNet(Encoder, num_features=2048, block_channel=[256, 512, 1024, 2048], parallel=parallel, pcamff=pcamff) elif model == 'hu': model = net.Hu(Encoder, num_features=2048, block_channel=[256, 512, 1024, 2048], semff=semff, pcamff=pcamff) else: raise NotImplementedError( "Select model type in [\'tbdp\', \'hu\']") if is_densenet: original_model = densenet.densenet161(pretrained=True) Encoder = modules.E_densenet(original_model) if model == 'tbdp': model = net.TBDPNet(Encoder, num_features=2208, block_channel=[192, 384, 1056, 2208], parallel=parallel, pcamff=pcamff) elif model == 'hu': model = net.Hu(Encoder, num_features=2208, block_channel=[192, 384, 1056, 2208], semff=semff, pcamff=pcamff) else: raise NotImplementedError( "Select model type in [\'tbdp\', \'hu\']") if is_senet: original_model = senet.senet154(pretrained='imagenet') Encoder = modules.E_senet(original_model) if model == 'tbdp': model = net.TBDPNet(Encoder, num_features=2048, block_channel=[256, 512, 1024, 2048], parallel=parallel, pcamff=pcamff) elif model == 'hu': model = net.Hu(Encoder, num_features=2048, block_channel=[256, 512, 1024, 2048], semff=semff, pcamff=pcamff) else: raise NotImplementedError( "Select model type in [\'tbdp\', \'hu\']") return model
import torch.nn.parallel import torch.nn.functional from utils import * from options import get_args from dataloader import nyudv2_dataloader from models.backbone_dict import backbone_dict from models import modules from models import net args = get_args('test') # lode nyud v2 test set TestImgLoader = nyudv2_dataloader.getTestingData_NYUDV2(args.batch_size, args.testlist_path, args.root_path) # model backbone = backbone_dict[args.backbone]() Encoder = modules.E_resnet(backbone) if args.backbone in ['resnet50']: model = net.model(Encoder, num_features=2048, block_channel=[256, 512, 1024, 2048], refinenet=args.refinenet) elif args.backbone in ['resnet18', 'resnet34']: model = net.model(Encoder, num_features=512, block_channel=[64, 128, 256, 512], refinenet=args.refinenet) model = nn.DataParallel(model).cuda() # load test model if args.loadckpt is not None and args.loadckpt.endswith('.pth.tar'): print("loading the specific model in checkpoint_dir: {}".format(args.loadckpt)) state_dict = torch.load(args.loadckpt) model.load_state_dict(state_dict) elif os.path.isdir(args.loadckpt): all_saved_ckpts = [ckpt for ckpt in os.listdir(args.loadckpt) if ckpt.endswith(".pth.tar")]