141 lines
3.6 KiB
JavaScript
141 lines
3.6 KiB
JavaScript
"use strict";
|
|
|
|
import path from 'path';
|
|
import fs from 'fs';
|
|
|
|
import brain from 'brain.js';
|
|
|
|
class AITrainer {
|
|
constructor({ ansi, settings, log, root_dir, GatewayRepo, DatasetFetcher }) {
|
|
this.a = ansi;
|
|
this.settings = settings;
|
|
this.root_dir = root_dir;
|
|
this.l = log;
|
|
this.dataset_fetcher = DatasetFetcher;
|
|
this.repo_gateway = GatewayRepo;
|
|
}
|
|
|
|
async train_unified() {
|
|
let filepath = path.join(
|
|
this.root_dir,
|
|
"..",
|
|
this.settings.ai.output_directory,
|
|
"ai.json"
|
|
);
|
|
|
|
if(!fs.existsSync(path.dirname(filepath)))
|
|
await fs.promises.mkdir(path.dirname(filepath), { recursive: true });
|
|
|
|
let training_result = await this.train_gateway(null, filepath);
|
|
|
|
let { latitude, longitude } = this.repo_gateway.get_average_location();
|
|
|
|
await fs.promises.writeFile(
|
|
path.join(
|
|
path.dirname(this.root_dir),
|
|
this.settings.ai.output_directory,
|
|
"index.json"
|
|
),
|
|
JSON.stringify({
|
|
properties: {
|
|
rssi_min: -150,
|
|
rssi_max: 0,
|
|
training_mode: "unified"
|
|
},
|
|
index: [{
|
|
id: "Unified AI (not a real gateway)",
|
|
filename: "ai.json",
|
|
latitude, longitude,
|
|
net_settings: training_result.net_settings
|
|
}]
|
|
}, null, "\t")
|
|
);
|
|
}
|
|
|
|
async train_all() {
|
|
let index = [];
|
|
for(let gateway of this.repo_gateway.iterate()) {
|
|
let filepath = path.join(this.root_dir, "..", this.settings.ai.output_directory, `${gateway.id}.json`);
|
|
|
|
if(!fs.existsSync(path.dirname(filepath)))
|
|
await fs.promises.mkdir(path.dirname(filepath), { recursive: true });
|
|
|
|
let result = await this.train_gateway(gateway.id, filepath);
|
|
if(!result || result.success === false) {
|
|
this.l.warn(`Warning: Failed to train AI for ${gateway.id}.`);
|
|
continue;
|
|
}
|
|
|
|
|
|
this.l.log(`Saved to ${filepath}`);
|
|
|
|
index.push({
|
|
id: gateway.id,
|
|
filename: path.basename(filepath),
|
|
latitude: gateway.latitude,
|
|
longitude: gateway.longitude,
|
|
net_settings: result.net_settings
|
|
});
|
|
}
|
|
|
|
await fs.promises.writeFile(
|
|
path.join(
|
|
path.dirname(this.root_dir),
|
|
this.settings.ai.output_directory,
|
|
"index.json"
|
|
),
|
|
JSON.stringify({
|
|
properties: {
|
|
rssi_min: -150,
|
|
rssi_max: 0,
|
|
training_mode: "split"
|
|
},
|
|
index
|
|
})
|
|
);
|
|
}
|
|
|
|
/**
|
|
* Trains an AI to predict the coverage of a specific gateway.
|
|
* @param {string} gateway_id The id of the gateway to train an AI for.
|
|
* @param {string} destination_filename The absolute path to the file to serialise the trained to. Required because we can't serialise and return a TensorFlow model, it has to be sent somewhere because the API is backwards and upside-down :-/
|
|
* @return {Promise} A promise that resolves when training and serialisation is complete.
|
|
*/
|
|
async train_gateway(gateway_id, destination_filename) {
|
|
this.l.log(`${this.a.fgreen}${this.a.hicol}Training AI for gateway ${gateway_id}${this.a.reset}`);
|
|
|
|
// Create the neural network
|
|
let net_settings = {
|
|
hiddenLayers: this.settings.ai.network_arch,
|
|
activation: "sigmoid"
|
|
};
|
|
let net = new brain.NeuralNetwork(net_settings);
|
|
|
|
// Fetch the dataset
|
|
let dataset = this.dataset_fetcher.fetch_all(gateway_id);
|
|
|
|
await net.trainAsync(dataset, {
|
|
iterations: this.settings.ai.epochs,
|
|
errorThresh: this.settings.ai.error_threshold,
|
|
|
|
learningRate: this.settings.ai.learning_rate,
|
|
momentum: this.settings.ai.momentum,
|
|
|
|
log: (log_line) => this.l.log(`[brain.js] ${log_line}`),
|
|
logPeriod: 50,
|
|
|
|
timeout: Infinity
|
|
});
|
|
|
|
await fs.promises.writeFile(destination_filename, JSON.stringify(net.toJSON()), null, "\t");
|
|
// console.log(result);
|
|
|
|
return {
|
|
success: true,
|
|
net_settings
|
|
};
|
|
}
|
|
}
|
|
|
|
export default AITrainer;
|