def configure_optimizers(self):
optimizer = object_from_dict(
self.hparams["optimizer"],
params=[x for x in self.model.parameters() if x.requires_grad],
)
scheduler = object_from_dict(self.hparams["scheduler"], optimizer=optimizer)
self.optimizers = [optimizer]
return self.optimizers, [scheduler]
def __init__(self, nSequence, resolution, cropSize, sequenceData): self.nSequence = nSequence self.resolution = resolution self.cropSize = cropSize shape_dict = dict([(key, value.shape) for key, value in sequenceData.frameData[0].__dict__.items()]) self.frameShape = utils.object_from_dict(shape_dict)
def main():
args = get_args()
torch.distributed.init_process_group(backend="nccl")
with open(args.config_path) as f:
hparams = yaml.load(f, Loader=yaml.SafeLoader)
hparams.update({
"local_rank": args.local_rank,
"fp16": args.fp16,
})
output_mask_path = args.output_path
output_mask_path.mkdir(parents=True, exist_ok=True)
hparams["output_mask_path"] = output_mask_path
device = torch.device("cuda", args.local_rank)
model = object_from_dict(hparams["model"])
model = model.to(device)
if args.fp16:
model = model.half()
corrections: Dict[str, str] = {"model.": ""}
state_dict = state_dict_from_disk(file_path=args.weight_path,
rename_in_layers=corrections)
model.load_state_dict(state_dict)
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[args.local_rank], output_device=args.local_rank)
file_paths = []
for regexp in ["*.jpg", "*.png", "*.jpeg", "*.JPG"]:
file_paths += sorted([x for x in tqdm(args.input_path.rglob(regexp))])
# Filter file paths for which we already have predictions
file_paths = [
x for x in file_paths
if not (args.output_path / x.parent.name / f"{x.stem}.png").exists()
]
dataset = InferenceDataset(file_paths,
transform=from_dict(hparams["test_aug"]))
sampler = DistributedSampler(dataset, shuffle=False)
dataloader = torch.utils.data.DataLoader(
dataset,
batch_size=args.batch_size,
num_workers=args.num_workers,
pin_memory=True,
shuffle=False,
drop_last=False,
sampler=sampler,
)
predict(dataloader, model, hparams, device)
def makeCudaTensors(self, x):
if isinstance(x, dict):
_dict = {key : self.makeCudaTensors(value) for key, value in x.items()}
return utils.object_from_dict(_dict)
elif isinstance(x, (list, tuple)):
return [self.makeCudaTensors(y) for y in x]
elif torch.is_tensor(x):
return x.cuda().float()
else:
return x
def __init__(self, hparams):
super().__init__()
self.hparams = hparams
self.model = object_from_dict(self.hparams["model"])
if "resume_from_checkpoint" in self.hparams:
corrections: Dict[str, str] = {"model.": ""}
state_dict = state_dict_from_disk(
file_path=self.hparams["resume_from_checkpoint"],
rename_in_layers=corrections,
)
self.model.load_state_dict(state_dict)
self.losses = [
("jaccard", 0.1, JaccardLoss(mode="binary", from_logits=True)),
("focal", 0.9, BinaryFocalLoss()),
]