def run_whole_training(experiment_name: str,
exp_config: ExperimenetConfig,
runs_dir="runs"):
model = get_model(exp_config.model_name, dropout=exp_config.dropout).cuda()
if exp_config.transfer_from_checkpoint:
transfer_checkpoint = fs.auto_file(exp_config.transfer_from_checkpoint)
print("Transferring weights from model checkpoint",
transfer_checkpoint)
checkpoint = load_checkpoint(transfer_checkpoint)
pretrained_dict = checkpoint["model_state_dict"]
transfer_weights(model, pretrained_dict)
if exp_config.resume_from_checkpoint:
checkpoint = load_checkpoint(
fs.auto_file(exp_config.resume_from_checkpoint))
unpack_checkpoint(checkpoint, model=model)
print("Loaded model weights from:", exp_config.resume_from_checkpoint)
report_checkpoint(checkpoint)
experiment_dir = os.path.join(runs_dir, experiment_name)
os.makedirs(experiment_dir, exist_ok=False)
config_fname = os.path.join(experiment_dir, f"config.json")
with open(config_fname, "w") as f:
f.write(json.dumps(jsonpickle.encode(exp_config), indent=2))
for stage in exp_config.stages:
run_stage_training(model,
stage,
exp_config,
experiment_dir=experiment_dir)
def ela_ecaresnext26tn_32x4d(num_classes=4, pretrained=True, dropout=0):
encoder = resnet.ecaresnext26tn_32x4d(in_chans=6)
del encoder.fc
if pretrained:
donor = resnet.ecaresnext26tn_32x4d(pretrained=True)
transfer_weights(encoder, donor.state_dict())
return TimmRgbElaModel(encoder, num_classes=num_classes, dropout=dropout)
def ela_skresnext50_32x4d(num_classes=4, pretrained=True, dropout=0):
encoder = skresnext50_32x4d(stem_type="deep", in_chans=6)
del encoder.fc
if pretrained:
donor = skresnext50_32x4d(pretrained=True)
transfer_weights(encoder, donor.state_dict())
return TimmRgbElaModel(encoder, num_classes=num_classes, dropout=dropout)
def ela_wider_resnet38(num_classes=4, pretrained=True, dropout=0):
from alaska2.models.backbones.wider_resnet import wider_resnet38
encoder = wider_resnet38(in_chans=6)
if pretrained:
checkpoint = torch.load(
fs.auto_file("wide_resnet38_ipabn_lr_256.pth.tar"),
map_location="cpu")
transfer_weights(encoder, checkpoint["state_dict"])
print("Loaded weights from Mapilary")
return TimmRgbElaModel(encoder, num_classes=num_classes, dropout=dropout)
def ela_tf_efficientnet_b2_ns(num_classes=4, pretrained=True, dropout=0):
encoder = efficientnet.tf_efficientnet_b2_ns(in_chans=9,
pretrained=False,
drop_path_rate=0.1)
del encoder.classifier
if pretrained:
donor = efficientnet.tf_efficientnet_b2_ns(pretrained=pretrained)
transfer_weights(encoder, donor.state_dict())
return TimmElaOnlyRichModel(encoder,
num_classes=num_classes,
dropout=dropout)
def ela_tf_efficientnet_b6_ns(num_classes=4, pretrained=True, dropout=0):
encoder = efficientnet.tf_efficientnet_b6_ns(in_chans=6,
pretrained=False,
drop_path_rate=0.2)
del encoder.classifier
if pretrained:
donor = efficientnet.tf_efficientnet_b6_ns(pretrained=pretrained)
transfer_weights(encoder, donor.state_dict())
return TimmRgbElaModel(
encoder,
num_classes=num_classes,
dropout=dropout,
mean=encoder.default_cfg["mean"],
std=encoder.default_cfg["std"],
)
def rgb_res_tf_efficientnet_b2_ns(num_classes=4, pretrained=True, dropout=0):
encoder = efficientnet.tf_efficientnet_b2_ns(in_chans=6,
pretrained=False,
drop_path_rate=0.1)
del encoder.classifier
if pretrained:
donor = efficientnet.tf_efficientnet_b2_ns(pretrained=pretrained)
transfer_weights(encoder, donor.state_dict())
return ImageAndResidualModel(
encoder,
num_classes=num_classes,
dropout=dropout,
mean=encoder.default_cfg["mean"],
std=encoder.default_cfg["std"],
)