55 lines
1.8 KiB
Python
55 lines
1.8 KiB
Python
"""Compact dual-head neural network (policy + value) sized for ESP32."""
|
|
|
|
from .config import CONV_FILTERS, NUM_CONV_LAYERS, DENSE_UNITS, LEARNING_RATE
|
|
|
|
|
|
def build_model():
|
|
"""Build a small AlphaZero-style network.
|
|
|
|
Input: (6, 7, 2) — current player pieces / opponent pieces
|
|
Output: policy (7,) — log-probabilities over columns
|
|
value (1,) — board evaluation in [-1, 1]
|
|
"""
|
|
from tensorflow import keras
|
|
from tensorflow.keras import layers
|
|
|
|
inp = layers.Input(shape=(6, 7, 2), name="board")
|
|
|
|
x = inp
|
|
for i in range(NUM_CONV_LAYERS):
|
|
x = layers.Conv2D(
|
|
CONV_FILTERS, 3, padding="same", activation="relu", name=f"conv{i}"
|
|
)(x)
|
|
x = layers.BatchNormalization(name=f"bn{i}")(x)
|
|
|
|
flat = layers.Flatten(name="flat")(x)
|
|
shared = layers.Dense(DENSE_UNITS, activation="relu", name="shared_dense")(flat)
|
|
|
|
# Policy head
|
|
policy = layers.Dense(7, name="policy_logits")(shared)
|
|
|
|
# Value head
|
|
value = layers.Dense(1, activation="tanh", name="value")(shared)
|
|
|
|
model = keras.Model(inputs=inp, outputs=[policy, value], name="connect4_net")
|
|
|
|
model.compile(
|
|
optimizer=keras.optimizers.Adam(learning_rate=LEARNING_RATE),
|
|
loss={
|
|
"policy_logits": keras.losses.CategoricalCrossentropy(from_logits=True),
|
|
"value": keras.losses.MeanSquaredError(),
|
|
},
|
|
loss_weights={"policy_logits": 1.0, "value": 1.0},
|
|
)
|
|
return model
|
|
|
|
|
|
def print_model_info(model):
|
|
model.summary()
|
|
total_params = model.count_params()
|
|
approx_size_kb = total_params * 4 / 1024 # float32
|
|
approx_int8_kb = total_params / 1024 # int8
|
|
print(f"\nTotal parameters: {total_params:,}")
|
|
print(f"Approx size (float32): {approx_size_kb:.1f} KB")
|
|
print(f"Approx size (int8): {approx_int8_kb:.1f} KB")
|