File size: 5,893 Bytes
bf928ee | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 | #!/usr/bin/env python3
"""Run a few in-memory training steps on DCU without modifying official weights."""
import argparse
import pickle
from pathlib import Path
from _bootstrap import DATA_DIR, WEIGHT_DIR
import numpy as np
import tensorflow as tf
import tf_keras
DEFAULT_DATA_DIR = DATA_DIR / "features" / "cdr_kmer3_ag_kmer1"
DEFAULT_WEIGHT_PATH = WEIGHT_DIR / "cdr_kmer3_ag_kmer1" / "Model99.h5"
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument("--data-dir", type=Path, default=DEFAULT_DATA_DIR)
parser.add_argument("--weights", type=Path, default=DEFAULT_WEIGHT_PATH)
parser.add_argument("--batch-size", type=int, default=8)
parser.add_argument("--steps", type=int, default=1)
return parser.parse_args()
def load_training_batch(data_dir, batch_size):
# Register features.InputFeatures before unpickling the official artifacts.
import model # noqa: F401
with (data_dir / "cdr_features_tr.pickle").open("rb") as reader:
cdr_features = pickle.load(reader)
with (data_dir / "ag_features_tr.pickle").open("rb") as reader:
ag_features = pickle.load(reader)
if len(cdr_features) != len(ag_features):
raise ValueError(
f"Unpaired training features: CDR={len(cdr_features)}, AG={len(ag_features)}"
)
if batch_size < 2 or batch_size > len(cdr_features):
raise ValueError(
f"--batch-size must be between 2 and {len(cdr_features)}, got {batch_size}"
)
selected_cdr = cdr_features[:batch_size]
selected_ag = ag_features[:batch_size]
inputs = [
np.asarray([item.input_ids for item in selected_cdr], dtype=np.int32),
np.asarray([item.cdr_number_ids for item in selected_cdr], dtype=np.int32),
np.asarray([item.input_ids for item in selected_ag], dtype=np.int32),
]
labels = np.asarray(
[[item.label_id] for item in selected_cdr], dtype=np.float32
)
expected_shapes = ((batch_size, 24), (batch_size, 24), (batch_size, 2371))
actual_shapes = tuple(array.shape for array in inputs)
if actual_shapes != expected_shapes:
raise ValueError(f"Unexpected input shapes: {actual_shapes}, expected {expected_shapes}")
if labels.shape != (batch_size, 1):
raise ValueError(f"Unexpected label shape: {labels.shape}")
return inputs, labels
def inference_loss(model, inputs, labels):
predictions = model(inputs, training=False)
losses = tf_keras.losses.binary_crossentropy(labels, predictions)
loss = float(tf.reduce_mean(losses).numpy())
predictions = predictions.numpy()
if not np.isfinite(loss) or not np.all(np.isfinite(predictions)):
raise FloatingPointError("Inference loss or predictions contain NaN/Inf")
return loss
def main():
args = parse_args()
if args.steps < 1:
raise ValueError(f"--steps must be at least 1, got {args.steps}")
if not args.data_dir.is_dir():
raise FileNotFoundError(f"Data directory not found: {args.data_dir}")
if not args.weights.is_file():
raise FileNotFoundError(f"Weight file not found: {args.weights}")
gpus = tf.config.list_physical_devices("GPU")
if not gpus:
raise RuntimeError(
"TensorFlow cannot see a DCU. Set HIP_VISIBLE_DEVICES and "
"CUDA_VISIBLE_DEVICES before starting Python."
)
for gpu in gpus:
tf.config.experimental.set_memory_growth(gpu, True)
inputs, labels = load_training_batch(args.data_dir, args.batch_size)
# Import after configuring the accelerator to avoid initializing it early.
from model import get_model
with tf.device("/GPU:0"):
model = get_model()
model.load_weights(str(args.weights))
loss_before = inference_loss(model, inputs, labels)
weights_before = [variable.numpy().copy() for variable in model.trainable_variables]
training_losses = []
for _ in range(args.steps):
loss = model.train_on_batch(inputs, labels)
training_losses.append(float(np.asarray(loss).reshape(-1)[0]))
loss_after = inference_loss(model, inputs, labels)
if not np.all(np.isfinite(training_losses)):
raise FloatingPointError(f"Training loss contains NaN/Inf: {training_losses}")
changed_count = 0
max_weight_change = 0.0
for before, variable in zip(weights_before, model.trainable_variables):
after = variable.numpy()
if not np.all(np.isfinite(after)):
raise FloatingPointError(f"Trainable variable contains NaN/Inf: {variable.name}")
change = float(np.max(np.abs(after - before)))
max_weight_change = max(max_weight_change, change)
if change > 0.0:
changed_count += 1
dcu_variable_count = sum("GPU:0" in variable.device for variable in model.variables)
print(f"TensorFlow: {tf.__version__}")
print(f"tf_keras: {tf_keras.__version__}")
print(f"DCU: {gpus[0]}")
print(f"Batch size: {args.batch_size}")
print(f"Steps: {args.steps}")
print(f"Labels: {labels.reshape(-1).tolist()}")
print(f"Inference loss before: {loss_before:.10g}")
print(f"Training losses: {training_losses}")
print(f"Inference loss after: {loss_after:.10g}")
print(f"Trainable tensors changed: {changed_count}/{len(weights_before)}")
print(f"Max absolute weight change: {max_weight_change:.10g}")
print(f"Variables placed on GPU:0: {dcu_variable_count}/{len(model.variables)}")
if changed_count == 0 or max_weight_change == 0.0:
raise AssertionError("No trainable weights changed after train_on_batch")
if dcu_variable_count == 0:
raise RuntimeError("No model variables were placed on GPU:0")
print("Minimal DCU training: PASS")
print("Official Model99.h5 was loaded read-only and was not overwritten.")
if __name__ == "__main__":
main()
|