def load_model() -> nn.Module: model = Unet(encoder_name="timm-efficientnet-b3", classes=1, encoder_weights=None) # TODO: While working on a fix, don't load for now :p state_dict = torch.load(MODEL_PATH)["state_dict"] state_dict = rename_layers(state_dict, {"model.": ""}) model.load_state_dict(state_dict) return model
def load_model_unet(_model_weights, is_inference=False): print('Using weights {}'.format(_model_weights)) if _model_weights == 'imagenet': model = Unet(unet_encoder, encoder_weights="imagenet", classes=4, activation=None, attention_type=ATTENTION_TYPE) if is_inference: model.eval() return model else: model = Unet( unet_encoder, encoder_weights=None, # "imagenet", classes=4, activation=None, attention_type=ATTENTION_TYPE) if is_inference: model.eval() if _model_weights is not None: device = torch.device("cuda") model.to(device) state = torch.load( _model_weights) # , map_location=lambda storage, loc: storage) model.load_state_dict(state["state_dict"]) optimizer_state = state['optimizer'] return model, optimizer_state # new_state_dict = OrderedDict() # # for k, v in state['state_dict'].items(): # if k in model.state_dict(): # new_state_dict[k] = v # model = model.load_state_dict(new_state_dict) return model
def load_models(model_path, config, n_gpus=1, unet=True): """ Loads segmentation models with a given architecture. """ # Load models models = glob(model_path + '/*fold_*.pth') models.sort() # List the models model_list = [] for fold in range(len(models)): if unet and n_gpus > 1: model = nn.DataParallel( Unet(config['model']['backbone'], encoder_weights="imagenet", activation='sigmoid')) elif unet: model = Unet(config['model']['backbone'], encoder_weights="imagenet", activation='sigmoid') elif n_gpus > 1: model = nn.DataParallel(EncoderDecoder(**config['model'])) else: model = EncoderDecoder(**config['model']) model.load_state_dict(torch.load(models[fold])) model_list.append(model) return model_list
testset = DataLoader( TestDataset(test_data_folder, df, mean, std), batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=True ) # Initialize mode and load trained weights ckpt_path = "../input/model_dump1/model.pth" device = torch.device("cuda") model = Unet("efficientnet-b5", encoder_weights=None, classes=4, activation=None) model.to(device) model.eval() state = torch.load(ckpt_path, map_location=lambda storage, loc: storage) model.load_state_dict(state["state_dict"]) # start prediction predictions = [] for i, batch in enumerate(tqdm(testset)): fnames, images = batch batch_preds = torch.sigmoid(model(images.to(device))) batch_preds = batch_preds.detach().cpu().numpy() for fname, preds in zip(fnames, batch_preds): for cls, pred in enumerate(preds): pred, num = post_process(pred, best_threshold[cls], min_size[cls]) rle = mask2rle(pred) name = fname + f"_{cls+1}" predictions.append([name, rle]) # save predictions to submission.csv
ENCODER = 'resnet34' ENCODER_WEIGHTS = 'imagenet' DEVICE = 'cuda' CLASSES = ['0', '1', '2', '3', '4'] ACTIVATION = 'softmax' unet_resnet34 = Unet( encoder_name=ENCODER, encoder_weights=None, classes=4, activation='sigmoid', ) state = torch.load("../input/bce-clf/unet_res34_525.pth") unet_resnet34.load_state_dict(state['model_state_dict']) unet_resnet34 = unet_resnet34.cuda() unet_resnet34 = unet_resnet34.eval() device = torch.device("cuda") model_senet = Unet('se_resnext50_32x4d', encoder_weights=None, classes=4, activation=None) model_senet.to(device) model_senet.eval() state = torch.load( '../input/senetmodels/senext50_30_epochs_high_threshold.pth', map_location=lambda storage, loc: storage) model_senet.load_state_dict(state["state_dict"])