Skip to content

Commit

Permalink
fix start in iterative scheduler, so the step will be calculated from…
Browse files Browse the repository at this point in the history
… this value
  • Loading branch information
Mikołaj Janusz committed Jun 10, 2024
1 parent 8081043 commit 7929f66
Showing 1 changed file with 10 additions and 3 deletions.
13 changes: 10 additions & 3 deletions pruning/architecture/pruning_methods/schedulers.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,16 +26,23 @@ def __iter__(self) -> Generator[list[float], None, None]:


class IterativeStepScheduler(BasePruningStepScheduler):
def __iter__(self) -> Generator[float, None, None]:
def __iter__(self) -> Generator[list[float], None, None]:
nonpruned_percent = 1

if self.start != 0:
yield [self.start]
nonpruned_percent -= round(self.start * nonpruned_percent, 8)

dummy_one = 1

# stop if pruned more than target pruning percentage - 0.1%
while nonpruned_percent - (1 - self.end) > 0.001:
current_step = round(self.step * nonpruned_percent, 8)
if self.start != 0:
current_step = round(self.step * dummy_one, 8)
dummy_one -= current_step
dummy_one = round(dummy_one, 8)
else:
current_step = round(self.step * nonpruned_percent, 8)

nonpruned_percent -= current_step
nonpruned_percent = round(nonpruned_percent, 8)

Expand Down

0 comments on commit 7929f66

Please sign in to comment.