Reduce peak memory usage when freezing parameters. (#14)

lkhphuc · web-flow · commit 1c6f5aa8adcf · 2022-08-16T14:01:20.000+02:00
diff --git a/big_vision/optax.py b/big_vision/optax.py
@@ -68,7 +68,7 @@ def create_schedule(mult=1.0, **kw):
       # Removes weight decay updates. Note that weight decay already has an
       # independent mask (which cannot be combined easily with a second mask),
       # so instead we multiply updates for frozen params with zero.
-      optax.masked(optax.scale(0.0), frozen_mask)
+      optax.masked(optax.set_to_zero(), frozen_mask)
   ]
 
   # Gradient clipping.

Original file line number	Diff line number	Diff line change
`@@ -68,7 +68,7 @@ def create_schedule(mult=1.0, **kw):`
`68`	`68`	`# Removes weight decay updates. Note that weight decay already has an`
`69`	`69`	`# independent mask (which cannot be combined easily with a second mask),`
`70`	`70`	`# so instead we multiply updates for frozen params with zero.`
`71`		`- optax.masked(optax.scale(0.0), frozen_mask)`
	`71`	`+ optax.masked(optax.set_to_zero(), frozen_mask)`
`72`	`72`	`]`
`73`	`73`
`74`	`74`	`# Gradient clipping.`