def create_model(device, type ='vgg16'): if type == 'vgg16': print('create vgg16 model') model = UNet16(pretrained=True) elif type == 'vgg16_bn': print('create vgg16_bn model') model = UNet16_bn(pretrained=True) elif type == 'vgg16_bn_do': print('create vgg16_bn_do model') model = UNet16_bn_do(pretrained=True) elif type == 'vgg16_fullbn_do': print('create vgg16_fullbn_do model') model = UNet16_fullbn_do(pretrained=True) elif type == 'resnet101': encoder_depth = 101 num_classes = 1 print('create resnet101 model') model = UNetResNet(encoder_depth=encoder_depth, num_classes=num_classes, pretrained=True) elif type == 'resnet34': encoder_depth = 34 num_classes = 1 print('create resnet34 model') model = UNetResNet(encoder_depth=encoder_depth, num_classes=num_classes, pretrained=True) elif type == 'unet+++': model = UNet_3Plus() print('create U-Net+++ model') elif type =='unet++': model = UNet_2Plus() print('create U-Net++ model') else: assert False model.eval() return model.to(device)
def create_model(device, type ='vgg16'): assert type == 'vgg16' or type == 'resnet101' if type == 'vgg16': model = UNet16(pretrained=True) elif type == 'resnet101': model = UNetResNet(pretrained=True, encoder_depth=101, num_classes=1) else: assert False model.eval() return model.to(device)
def create_model(device, type='vgg16'): assert type == 'vgg16' or type == 'resnet101' if type == 'vgg16': print('create vgg16 model') model = UNet16(pretrained=True) elif type == 'resnet101': encoder_depth = 101 num_classes = 1 print('create resnet101 model') model = UNetResNet(encoder_depth=encoder_depth, num_classes=num_classes, pretrained=True) else: assert False model.eval() return model.to(device)