misc fixes
This commit is contained in:
+312
-132
@@ -1,134 +1,314 @@
|
||||
// use crate::model::{ChessBatcher, ChessModel, ChessModelConfig};
|
||||
// use burn::config::Config;
|
||||
// use burn::data::dataloader::batcher::Batcher;
|
||||
// use burn::optim::{Adam, AdamConfig, Optimizer, SimpleOptimizer};
|
||||
// use burn::optim::adaptor::OptimizerAdaptor;
|
||||
// use burn::prelude::Backend;
|
||||
// use burn::tensor::backend::AutodiffBackend;
|
||||
// use burn::train::TrainStep;
|
||||
// use crate::mcts::MctsResults;
|
||||
//
|
||||
// #[derive(Config, Debug)]
|
||||
// pub struct ChessTrainerConfig {
|
||||
// pub model: ChessModelConfig,
|
||||
// pub optimizer: AdamConfig,
|
||||
// #[config(default = 10)]
|
||||
// pub num_epochs: usize,
|
||||
// #[config(default = 64)]
|
||||
// pub batch_size: usize,
|
||||
// #[config(default = 4)]
|
||||
// pub num_workers: usize,
|
||||
// #[config(default = 42)]
|
||||
// pub seed: u64,
|
||||
// #[config(default = 1.0e-4)]
|
||||
// pub learning_rate: f64,
|
||||
// }
|
||||
//
|
||||
// impl ChessTrainerConfig {
|
||||
// pub fn init<B: Backend>(
|
||||
// model_config: ChessModelConfig,
|
||||
// optimizer: AdamConfig,
|
||||
// num_epochs: usize,
|
||||
// batch_size: usize,
|
||||
// num_workers: usize,
|
||||
// seed: u64,
|
||||
// learning_rate: f64,
|
||||
// ) -> ChessTrainer<B> {
|
||||
// ChessTrainer {
|
||||
// model: model_config::init(),
|
||||
// optimizer: optimizer.init(),
|
||||
// num_epochs,
|
||||
//
|
||||
//
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
//
|
||||
// pub struct ChessTrainer<B: AutodiffBackend> {
|
||||
// pub model: ChessModel<B>,
|
||||
// pub optimizer: Adam,
|
||||
// learning_rate: f64,
|
||||
// pub batcher: ChessBatcher,
|
||||
// pub device: B::Device,
|
||||
// }
|
||||
//
|
||||
// impl<B: AutodiffBackend> ChessTrainer<B> {
|
||||
// pub fn new(model: ChessModel<B>, device: B::Device) -> Self {
|
||||
// let optimizer = AdamConfig::new().init();
|
||||
//
|
||||
// Self {
|
||||
// model,
|
||||
// optimizer,
|
||||
// batcher: ChessBatcher::default(),
|
||||
// device,
|
||||
// }
|
||||
// }
|
||||
//
|
||||
// pub fn train_step(&mut self, batch_data: Vec<MctsResults>) -> f32 {
|
||||
// // 1. Convert to tensors
|
||||
// let batch = self.batcher.batch(batch_data, &self.device);
|
||||
//
|
||||
// // 2. Forward + backward (your TrainStep impl)
|
||||
// let output = self.model.step(batch);
|
||||
//
|
||||
// // 3. Update weights
|
||||
// self.optimizer.step(, &mut self.model, output.grads);
|
||||
//
|
||||
// // 4. Return loss (for logging)
|
||||
// let loss_tensor = output.item.loss.clone().into_data();
|
||||
// let loss = loss_tensor.to_vec::<f32>().unwrap()[0];
|
||||
//
|
||||
// loss
|
||||
// }
|
||||
// }
|
||||
use crate::mcts::{BoardState, BoardStateStatus, Mcts, MctsConfig, MctsResults};
|
||||
use crate::net::encoding::decode_move;
|
||||
use crate::net::model::{
|
||||
ChessBatcher, ChessModel, ChessModelConfig, ModelMetadata, TrainingSample,
|
||||
};
|
||||
use burn::data::dataloader::batcher::Batcher;
|
||||
use burn::module::{AutodiffModule, Module};
|
||||
use burn::optim::{AdamConfig, GradientsParams, Optimizer};
|
||||
// use burn::record::{FullPrecisionSettings, NamedMpkFileRecorder};
|
||||
use burn::prelude::Device;
|
||||
use burn::store::ModuleRecord;
|
||||
use chess::{ChessMove, Color};
|
||||
use rand::rngs::SmallRng;
|
||||
use rand::seq::SliceRandom;
|
||||
use rand::{RngExt, SeedableRng};
|
||||
use std::collections::VecDeque;
|
||||
use std::fs::File;
|
||||
use std::io::BufReader;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::Arc;
|
||||
use std::time::Instant;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
// fn create_artifact_dir(artifact_dir: &str) {
|
||||
// // Remove existing artifacts before to get an accurate learner summary
|
||||
// std::fs::remove_dir_all(artifact_dir).ok();
|
||||
// std::fs::create_dir_all(artifact_dir).ok();
|
||||
// }
|
||||
//
|
||||
// pub fn train<B: AutodiffBackend>(
|
||||
// artifact_dir: &str,
|
||||
// config: ChessTrainingConfig,
|
||||
// device: B::Device,
|
||||
// ) {
|
||||
// create_artifact_dir(artifact_dir);
|
||||
// config
|
||||
// .save(format!("{artifact_dir}/config.json"))
|
||||
// .expect("Config should be saved successfully");
|
||||
//
|
||||
// B::seed(&device, config.seed);
|
||||
//
|
||||
// let batcher = ChessBatcher::default();
|
||||
pub struct TrainingConfig {
|
||||
pub max_time_s: Option<u64>,
|
||||
pub num_iters: Option<usize>,
|
||||
pub max_depth: u16, // unused
|
||||
pub model_name: String,
|
||||
pub load_model: bool,
|
||||
pub hidden_channels: usize,
|
||||
pub num_blocks: usize,
|
||||
pub batch_size: usize, // num positions sampled to a batch (2048)
|
||||
pub num_episodes: usize, // num games generated per iteration (5000)
|
||||
pub buffer_max_size: usize, // max number of samples in the buffer
|
||||
pub mcts_config: MctsConfig,
|
||||
pub optimizer: AdamConfig,
|
||||
pub lr: f64,
|
||||
pub seed: Option<u64>,
|
||||
}
|
||||
|
||||
// let dataloader_train = DataLoaderBuilder::new(batcher.clone())
|
||||
// .batch_size(config.batch_size)
|
||||
// .shuffle(config.seed)
|
||||
// .num_workers(config.num_workers)
|
||||
// .build(MnistDataset::train());
|
||||
//
|
||||
// let dataloader_test = DataLoaderBuilder::new(batcher)
|
||||
// .batch_size(config.batch_size)
|
||||
// .shuffle(config.seed)
|
||||
// .num_workers(config.num_workers)
|
||||
// .build(MnistDataset::test());
|
||||
//
|
||||
// let training = SupervisedTraining::new(artifact_dir, dataloader_train, dataloader_test)
|
||||
// .metrics((AccuracyMetric::new(), LossMetric::new()))
|
||||
// .with_file_checkpointer(CompactRecorder::new())
|
||||
// .num_epochs(config.num_epochs)
|
||||
// .summary();
|
||||
//
|
||||
// let model = config.model.init::<B>(&device);
|
||||
// let result = training.launch(Learner::new(
|
||||
// model,
|
||||
// config.optimizer.init(),
|
||||
// config.learning_rate,
|
||||
// ));
|
||||
//
|
||||
// result
|
||||
// .model
|
||||
// .save_file(format!("{artifact_dir}/model"), &CompactRecorder::new())
|
||||
// .expect("Trained model should be saved successfully");
|
||||
// }
|
||||
pub fn train(training_config: TrainingConfig, device: &Device) {
|
||||
let model_path = format!("artifacts/{}", training_config.model_name.as_str());
|
||||
let metadata_path = format!("{}.json", model_path);
|
||||
println!("Creating model...");
|
||||
let mut model: ChessModel = ChessModelConfig::init(
|
||||
training_config.num_blocks,
|
||||
training_config.hidden_channels,
|
||||
device,
|
||||
);
|
||||
if training_config.load_model {
|
||||
println!("Loading model {}...", model_path);
|
||||
// Load model in full precision from MessagePack file
|
||||
let record = ModuleRecord::load(&model_path).expect("Failed to load model.");
|
||||
model = model.load_record(record)
|
||||
}
|
||||
|
||||
let train = Arc::new(AtomicBool::new(true));
|
||||
let train_signal = Arc::clone(&train);
|
||||
|
||||
let mut iter: usize = 0;
|
||||
|
||||
if training_config.load_model {
|
||||
let file = File::open(&metadata_path).unwrap();
|
||||
|
||||
// 2. Use a buffered reader for efficiency
|
||||
let reader = BufReader::new(file);
|
||||
|
||||
// 3. Deserialize JSON directly from the file reader
|
||||
let metadata: ModelMetadata = serde_json::from_reader(reader).unwrap();
|
||||
iter = metadata.iterations;
|
||||
println!("Loaded model had {} iters", iter);
|
||||
}
|
||||
|
||||
let start_time = Instant::now();
|
||||
|
||||
ctrlc::set_handler(move || {
|
||||
train_signal.store(false, Ordering::Relaxed);
|
||||
println!("Finishing batch before exiting...");
|
||||
})
|
||||
.expect("Error setting Ctrl-C handler");
|
||||
|
||||
let mut replay_buffer: VecDeque<TrainingSample> = VecDeque::new();
|
||||
|
||||
let mut mcts: Mcts = Mcts {
|
||||
config: training_config.mcts_config,
|
||||
};
|
||||
|
||||
// Create RNG once and reuse it for sampling and shuffling
|
||||
// Seed from system time (platform default entropy may be unavailable in some contexts)
|
||||
let now = SystemTime::now().duration_since(UNIX_EPOCH).unwrap();
|
||||
let seed = training_config.seed.unwrap_or(now.as_nanos() as u64);
|
||||
println!("using seed: {}", seed);
|
||||
let mut rng = SmallRng::seed_from_u64(seed);
|
||||
|
||||
// Initialize optimizer once so state (moments) persist across steps
|
||||
let mut optim = training_config.optimizer.init();
|
||||
|
||||
println!("Starting training...");
|
||||
while train.load(Ordering::Relaxed) {
|
||||
println!("Iteration: {}", iter);
|
||||
let infer_model = model.valid();
|
||||
// Gen samples
|
||||
println!("Generating {} games...", training_config.num_episodes);
|
||||
for episode in 0..training_config.num_episodes {
|
||||
println!("Episode: {}", episode);
|
||||
let mut board_state = BoardState::default();
|
||||
let mut episode_buffer: Vec<MctsResults> = vec![];
|
||||
|
||||
let mut game_hist = String::new();
|
||||
|
||||
let mut ply_count: usize = 0;
|
||||
|
||||
while board_state.status == BoardStateStatus::Ongoing {
|
||||
// let before = Instant::now();
|
||||
let results = mcts.search(&board_state, &infer_model, &device);
|
||||
// println!("{}", before.elapsed().as_millis());
|
||||
episode_buffer.push(results);
|
||||
|
||||
let temp = if ply_count < 30 {
|
||||
1.0
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
|
||||
let adjusted = apply_temperature(&episode_buffer.last().unwrap().move_dist, temp);
|
||||
|
||||
let stm = board_state.board.side_to_move();
|
||||
let mv = sample_move(&adjusted, &mut rng, stm).unwrap();
|
||||
// println!("playing move: {}", mv);
|
||||
if episode == 0 {
|
||||
game_hist.push_str(mv.to_string().as_str());
|
||||
game_hist.push(' ');
|
||||
}
|
||||
|
||||
board_state.apply_move(mv);
|
||||
ply_count += 1;
|
||||
if ply_count >= 400 && board_state.status == BoardStateStatus::Ongoing {
|
||||
board_state.status = BoardStateStatus::FiftyMove;
|
||||
}
|
||||
}
|
||||
if episode == 0 {
|
||||
println!("Game history of first game of iteration: {}", game_hist);
|
||||
}
|
||||
|
||||
for result in episode_buffer.iter() {
|
||||
let stm = result.board_state.board.side_to_move();
|
||||
let outcome = match board_state.status {
|
||||
BoardStateStatus::WhiteWinner => {
|
||||
if stm == Color::White {
|
||||
1.0
|
||||
} else {
|
||||
-1.0
|
||||
}
|
||||
}
|
||||
BoardStateStatus::BlackWinner => {
|
||||
if stm == Color::Black {
|
||||
1.0
|
||||
} else {
|
||||
-1.0
|
||||
}
|
||||
}
|
||||
_ => 0.0,
|
||||
};
|
||||
|
||||
replay_buffer.push_back(TrainingSample::from_mcts_with_outcome(
|
||||
result.clone(),
|
||||
outcome,
|
||||
));
|
||||
|
||||
if replay_buffer.len() > training_config.buffer_max_size {
|
||||
replay_buffer.pop_front();
|
||||
}
|
||||
}
|
||||
}
|
||||
// train
|
||||
|
||||
if replay_buffer.is_empty() {
|
||||
println!("Replay buffer is empty, skipping training step.");
|
||||
continue;
|
||||
}
|
||||
|
||||
let num_samples = std::cmp::min(training_config.batch_size, replay_buffer.len());
|
||||
println!(
|
||||
"Finished self-play! Training on {} samples (buffer size: {})...",
|
||||
num_samples,
|
||||
replay_buffer.len()
|
||||
);
|
||||
let mut indices: Vec<usize> = (0..replay_buffer.len()).collect();
|
||||
indices.shuffle(&mut rng);
|
||||
|
||||
let samples = indices
|
||||
.into_iter()
|
||||
.take(num_samples)
|
||||
.map(|i| replay_buffer[i].clone())
|
||||
.collect::<Vec<TrainingSample>>();
|
||||
|
||||
let batcher = ChessBatcher {};
|
||||
|
||||
let batch = batcher.batch(samples, &device);
|
||||
|
||||
let output = model.forward_chess(batch.states, batch.policy_targets, batch.value_targets);
|
||||
|
||||
let loss_val = output.loss.clone().into_data().to_vec::<f32>().unwrap()[0];
|
||||
let p_loss_val = output.policy_loss.clone().into_data().to_vec::<f32>().unwrap()[0];
|
||||
let v_loss_val = output.value_loss.clone().into_data().to_vec::<f32>().unwrap()[0];
|
||||
|
||||
println!(
|
||||
"Iteration {} loss | Total: {:.4} | Policy: {:.4} | Value: {:.4}",
|
||||
iter, loss_val, p_loss_val, v_loss_val
|
||||
);
|
||||
|
||||
let grads = output.loss.backward();
|
||||
let grads = GradientsParams::from_grads(grads, &model);
|
||||
|
||||
model = optim.step(training_config.lr, model, grads);
|
||||
|
||||
iter += 1;
|
||||
|
||||
if iter % 100 == 0 {
|
||||
println!("Completed {} iterations", iter);
|
||||
}
|
||||
|
||||
if training_config.max_time_s.is_some()
|
||||
&& start_time.elapsed().as_secs() > training_config.max_time_s.unwrap()
|
||||
{
|
||||
println!("Training stopping due to time limit...");
|
||||
break;
|
||||
}
|
||||
|
||||
if training_config.num_iters.is_some() && iter >= training_config.num_iters.unwrap() {
|
||||
println!(
|
||||
"Training stopping due to iteration limit ({} iters completed)...",
|
||||
iter
|
||||
);
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
println!("Saving model...");
|
||||
|
||||
let metadata = ModelMetadata {
|
||||
name: training_config.model_name,
|
||||
iterations: iter,
|
||||
};
|
||||
|
||||
std::fs::write(
|
||||
metadata_path,
|
||||
serde_json::to_string_pretty(&metadata)
|
||||
.expect("Should be able to convert metadata to JSON string"),
|
||||
)
|
||||
.expect("Should be able to write metadata");
|
||||
|
||||
model
|
||||
.into_record()
|
||||
.save(&model_path)
|
||||
.expect("Should be able to save the model");
|
||||
println!("Model saved in {:?}, exiting training.", model_path);
|
||||
return;
|
||||
}
|
||||
|
||||
fn apply_temperature(visits: &[(usize, f32)], temperature: f32) -> Vec<(usize, f32)> {
|
||||
if visits.is_empty() {
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
// Special case: deterministic selection
|
||||
if temperature == 0.0 {
|
||||
let (&best_idx, _) = visits
|
||||
.iter()
|
||||
.max_by(|a, b| a.1.partial_cmp(&b.1).unwrap())
|
||||
.map(|(i, p)| (i, p))
|
||||
.unwrap();
|
||||
|
||||
return vec![(best_idx, 1.0)];
|
||||
}
|
||||
|
||||
let inv_temp = 1.0 / temperature;
|
||||
|
||||
// Step 1: apply exponent
|
||||
let mut adjusted: Vec<(usize, f32)> =
|
||||
visits.iter().map(|(i, v)| (*i, v.powf(inv_temp))).collect();
|
||||
|
||||
// Step 2: normalize
|
||||
let sum: f32 = adjusted.iter().map(|(_, v)| *v).sum();
|
||||
|
||||
if sum <= 0.0 {
|
||||
return adjusted; // fallback (shouldn't happen in normal MCTS)
|
||||
}
|
||||
|
||||
for (_, v) in adjusted.iter_mut() {
|
||||
*v /= sum;
|
||||
}
|
||||
|
||||
adjusted
|
||||
}
|
||||
|
||||
fn sample_move(
|
||||
dist: &[(usize, f32)],
|
||||
rng: &mut SmallRng,
|
||||
side_to_move: Color,
|
||||
) -> Option<ChessMove> {
|
||||
let mut r: f32 = rng.random_range(0.0..1.0);
|
||||
|
||||
for (idx, p) in dist {
|
||||
r -= *p;
|
||||
if r <= 0.0 {
|
||||
return Some(decode_move(*idx, side_to_move).expect("Invalid move"));
|
||||
}
|
||||
}
|
||||
|
||||
// fallback due to floating point drift
|
||||
dist.get(0)
|
||||
.map(|(idx, _)| decode_move(*idx, side_to_move).expect("Invalid move"))
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user