sbrl-github/research-rainfallradar
1
0
Fork 0
mirror of https://github.com/sbrl/research-rainfallradar synced 2024-12-22 22:25:01 +00:00
Code Issues Projects Releases Packages Wiki Activity

dlr eo: add JIT_COMPILE and MIXED_PRECISION

Browse source
This commit is contained in:
Starbeamrainbowlabs 2023-02-02 16:14:09 +00:00
parent 71088b8c0b
commit 2bf1872aca
Signed by: sbrl
GPG key ID: 1BE5172E637709C2
1 changed files with 10 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

13
aimodel/src/encoderonly_test_rainfall.py
View file

@ -28,6 +28,8 @@ WINDOW_SIZE = int(os.environ["WINDOW_SIZE"]) if "WINDOW_SIZE" in os.environ e
STEPS_PER_EPOCH = int(os.environ["STEPS_PER_EPOCH"]) if "STEPS_PER_EPOCH" in os.environ else None
STEPS_PER_EXECUTION = int(os.environ["STEPS_PER_EXECUTION"]) if "STEPS_PER_EXECUTION" in os.environ else None
LEARNING_RATE = float(os.environ["LEARNING_RATE"]) if "LEARNING_RATE" in os.environ else 0.001
JIT_COMPILE = True if "JIT_COMPILE" in os.environ else False
MIXED_PRECISION = True if "MIXED_PRECISION" in os.environ else False
logger.info("Encoder-only rainfall radar TEST")
logger.info(f"> DIRPATH_RAINFALLWATER {DIRPATH_RAINFALLWATER}")
@ -43,6 +45,9 @@ logger.info(f"> LEARNING_RATE {LEARNING_RATE}")
if not os.path.exists(DIRPATH_OUTPUT):
os.makedirs(os.path.join(DIRPATH_OUTPUT, "checkpoints"))
if MIXED_PRECISION:
tf.keras.mixed_precision.set_policy("mixed_float16")
# ██████ █████ ████████ █████ ███████ ███████ ████████
# ██ ██ ██ ██ ██ ██ ██ ██ ██ ██
@ -65,7 +70,7 @@ dataset_train, dataset_validate = dataset_encoderonly(
# ██ ██ ██ ██ ██ ██ ██ ██ ██
# ██ ██ ██████ ██████ ███████ ███████
def make_encoderonly(windowsize, channels, encoder="convnext", water_bins=2, steps_per_execution=1, **kwargs):
def make_encoderonly(windowsize, channels, encoder="convnext", water_bins=2, steps_per_execution=1, jit_compile=False, **kwargs):
if encoder == "convnext":
model = make_convnext(
input_shape=(windowsize, windowsize, channels),
@ -97,7 +102,8 @@ def make_encoderonly(windowsize, channels, encoder="convnext", water_bins=2, ste
metrics = [
tf.keras.metrics.SparseCategoricalAccuracy()
],
steps_per_execution=steps_per_execution
steps_per_execution=steps_per_execution,
jit_compile=jit_compile
)
return model
@ -106,7 +112,8 @@ def make_encoderonly(windowsize, channels, encoder="convnext", water_bins=2, ste
model = make_encoderonly(
windowsize=WINDOW_SIZE,
channels=CHANNELS,
steps_per_execution=STEPS_PER_EXECUTION
steps_per_execution=STEPS_PER_EXECUTION,
jit_compile=JIT_COMPILE
)
summarywriter(model, os.path.join(DIRPATH_OUTPUT, "summary.txt"))