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"])