LoRaWAN-Signal-Mapping/server/train-ai/AITrainer.mjs

99 lines
2.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;
}
generate_neural_net() {
let net = new brain.NeuralNetwork({
hiddenLayers: this.settings.ai.network_arch,
activation: "sigmoid"
});
return net;
}
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`);
console.log(filepath);
if(!fs.existsSync(path.dirname(filepath)))
await fs.promises.mkdir(path.dirname(filepath), { recursive: true });
if(!await this.train_gateway(gateway.id, filepath)) {
this.l.warn(`Warning: Failed to train AI for ${gateway.id}.`);
continue;
}
index.push({
id: gateway.id,
filename: path.basename(filepath),
latitude: gateway.latitude,
longitude: gateway.longitude
});
}
await fs.promises.writeFile(
path.join(
path.dirname(this.root_dir),
this.settings.ai.output_directory,
"index.json"
),
JSON.stringify({
properties: {
rssi_min: this.settings.ai.rssi_min,
rssi_max: this.settings.ai.rssi_max
},
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}`);
let net = this.generate_neural_net();
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: true,
logPeriod: 50,
timeout: Infinity
});
await fs.promises.writeFile(destination_filename, JSON.stringify(net.toJSON()));
// console.log(result);
return true;
}
}
export default AITrainer;