Support multiple click interactions per batch for DeepEdit training
@@ -37,6 +37,7 @@ class Interaction:
train: True for training mode or False for evaluation mode
click_probability_key: key to click/interaction probability
label_names: Dict of label names
+ max_interactions: maximum number of interactions per iteration
"""
def __init__(
@@ -44,8 +45,9 @@ class Interaction:
deepgrow_probability: float,
transforms: Union[Sequence[Callable], Callable],
train: bool,
- label_names: Dict[str, int],
+ label_names: Union[None, Dict[str, int]] = None,
click_probability_key: str = "probability",
+ max_interactions: int = 1,
) -> None:
self.deepgrow_probability = deepgrow_probability
@@ -53,21 +55,20 @@ class Interaction:
self.train = train
self.label_names = label_names
self.click_probability_key = click_probability_key
+ self.max_interactions = max_interactions
def __call__(self, engine: Union[SupervisedTrainer, SupervisedEvaluator], batchdata: Dict[str, torch.Tensor]):
-
if batchdata is None:
raise ValueError("Must provide batch data for current iteration.")
if np.random.choice([True, False], p=[self.deepgrow_probability, 1 - self.deepgrow_probability]):
-
- # Run the inner loop only once
+ for j in range(self.max_interactions):
inputs, _ = engine.prepare_batch(batchdata)
inputs = inputs.to(engine.state.device)
engine.fire_event(IterationEvents.INNER_ITERATION_STARTED)
-
engine.network.eval()
+
with torch.no_grad():
if engine.amp:
with torch.cuda.amp.autocast():
@@ -78,15 +79,14 @@ class Interaction:
# decollate/collate batchdata to execute click transforms
batchdata_list = decollate_batch(batchdata, detach=True)
-
for i in range(len(batchdata_list)):
- batchdata_list[i][self.click_probability_key] = 1.0
+ batchdata_list[i][self.click_probability_key] = (
+ (1.0 - ((1.0 / self.max_interactions) * j)) if self.train else 1.0
+ )
batchdata_list[i] = self.transforms(batchdata_list[i])
batchdata = list_data_collate(batchdata_list)
-
engine.fire_event(IterationEvents.INNER_ITERATION_COMPLETED)
-
else:
# zero out input guidance channels
batchdata_list = decollate_batch(batchdata, detach=True)
@@ -96,5 +96,4 @@ class Interaction:
# first item in batch only
engine.state.batch = batchdata
-
return engine._iteration(engine, batchdata)
