diff --git a/examples/src/gpu_autoencoder_example.js b/examples/src/gpu_autoencoder_example.js index 1fb5f6f..42f3e94 100644 --- a/examples/src/gpu_autoencoder_example.js +++ b/examples/src/gpu_autoencoder_example.js @@ -20,7 +20,8 @@ import { Dense, ChainModel, ProgressUtils, - Iter, Matrix, ImageUtils + ImageUtils, + Matrix } from "mind-net.js"; import {GpuModelWrapper} from "@mind-net.js/gpu" @@ -31,138 +32,157 @@ import * as ModelUtils from "./utils/model.js"; console.log("Fetching datasets..."); -//const DatasetBigUrl = "https://github.com/DrA1ex/mind-net.js/files/12456697/mnist-10000-28.zip"; -//const DatasetBigUrl = "https://github.com/DrA1ex/mind-net.js/files/12398103/cartoon-2500-64.zip"; -const DatasetBigUrl = "https://github.com/DrA1ex/mind-net.js/files/12407792/cartoon-2500-28.zip"; - -const zipData = await ProgressUtils.fetchProgress(DatasetBigUrl); +const DatasetSmallUrl = "https://github.com/DrA1ex/mind-net.js/files/12684744/cartoon_avatar_1000_32.zip"; +const DatasetBigUrl = "https://github.com/DrA1ex/mind-net.js/files/12684745/cartoon_avatar_1000_64.zip"; +const TestDatasetUrl = "https://github.com/DrA1ex/mind-net.js/files/12684746/face-25-32.zip"; console.log("Preparing...") -const zipLoadingProgress = ProgressUtils.progressCallback({ +const getLoadingProgressFn = () => ProgressUtils.progressCallback({ update: true, color: ProgressUtils.Color.magenta, limit: ProgressUtils.ValueLimit.inclusive, }); -const loadedSet = ImageUtils.grayscaleDataset(await DatasetUtils.loadDataset(zipData.buffer, zipLoadingProgress)); - -// Create variations with different positioning -const setMul = 3; -const dataSet = new Array(loadedSet.length * setMul); -for (let k = 0; k < setMul; k++) { - for (let i = 0; i < loadedSet.length; i++) { - const [x, y] = [Math.random() * 2 - 1, Math.random() * 2 - 1]; - dataSet[k * loadedSet.length + i] = ImageUtils.shiftImage(loadedSet[i], Math.round(x), Math.round(y), 1); - } -} -const trainIterations = 100; -const epochsPerIter = 5; -const batchSize = 128; +const CNT = 2500; + +let smallZipData = await ProgressUtils.fetchProgress(DatasetSmallUrl); +const smallDataSet = ImageUtils.grayscaleDataset( + (await DatasetUtils.loadDataset(smallZipData.buffer, getLoadingProgressFn())).slice(0, CNT) +); +smallZipData = undefined; + +let bigZipData = await ProgressUtils.fetchProgress(DatasetBigUrl); +const bigDataSet = ImageUtils.grayscaleDataset( + (await DatasetUtils.loadDataset(bigZipData.buffer, getLoadingProgressFn())).slice(0, CNT) +); +bigZipData = undefined; + +let testZipData = await ProgressUtils.fetchProgress(TestDatasetUrl); +const testColorfulDataSet = await DatasetUtils.loadDataset(testZipData.buffer, getLoadingProgressFn()); +testZipData = undefined; + +const trainIterations = 50; +const epochsPerIterFilter = 10; +const epochsPerIterUpscaler = 5; +const batchSize = 256; +const batchSizeUpscaler = 512; const imageChannel = 1; -const sampleScale = 4; -const epochSampleSize = 10; -const finalSampleSize = 20; +const testImageChannel = 3; +const epochSampleSize = Math.min(10, Math.floor(Math.sqrt(testColorfulDataSet.length))); +const finalSampleSize = Math.min(20, Math.floor(Math.sqrt(testColorfulDataSet.length))); const outPath = "./out"; -const lr = 0.001; +const lr = 0.0003; +const lrUpscaler = 0.0005; const decay = 5e-4; const beta1 = 0.5; -const dropout = 0.3; -const activation = "leakyRelu"; -const loss = "mse"; +const dropout = 0.2; +const activation = "relu"; +const loss = "l2"; const initializer = "xavier"; -const LatentSpaceSize = 32; -const Sizes = [256, 128, 64]; - -const encoder = new SequentialModel(new AdamOptimizer({lr, decay, beta1}), loss); -encoder.addLayer(new Dense(dataSet[0].length)); -for (const size of Sizes) { - encoder.addLayer( - new Dense(size, { - activation, - weightInitializer: initializer, - options: {dropout} - }) - ); +const FilterSizes = [512, 384, 256, 384, 512]; +const UpscalerSizes = [512, 640, 784, 896, 1024]; + +function createOptimizer(lr) { + return new AdamOptimizer({lr, decay, beta1}); } -encoder.addLayer(new Dense(LatentSpaceSize, {activation, weightInitializer: initializer})); - -const decoder = new SequentialModel(new AdamOptimizer({lr, decay, beta1}), loss); -decoder.addLayer(new Dense(LatentSpaceSize)); -for (const size of Iter.reverse(Sizes)) { - decoder.addLayer( - new Dense(size, { - activation, - weightInitializer: initializer, - options: {dropout} - }) - ); + +function createHiddenLayer(size) { + return new Dense(size, { + activation, + weightInitializer: initializer, + options: {dropout} + }); } -decoder.addLayer(new Dense(dataSet[0].length, {weightInitializer: initializer, activation: "tanh"})); -const chain = new ChainModel(new AdamOptimizer({lr, decay, beta1}), loss); -chain.addModel(encoder); -chain.addModel(decoder); +function createModel(inSize, outSize, hiddenLayers, lr) { + const model = new SequentialModel(createOptimizer(lr), loss); + model.addLayer(new Dense(inSize)); + for (const size of hiddenLayers) model.addLayer(createHiddenLayer(size)); + model.addLayer(new Dense(outSize, {activation: "tanh"})); + model.compile(); -chain.compile(); + return model; +} -const pEncoder = new GpuModelWrapper(encoder, {batchSize}); -const pDecoder = new GpuModelWrapper(decoder, {batchSize}); -const pChain = new GpuModelWrapper(chain, {batchSize}); +const filterModel = createModel(smallDataSet[0].length, smallDataSet[0].length, FilterSizes, lr); +const upscalerModel = createModel(smallDataSet[0].length, bigDataSet[0].length, UpscalerSizes, lrUpscaler); -async function _decode(from, to) { - const encFrom = pEncoder.compute(from); - const encTo = pEncoder.compute(to); +const resultingModel = new ChainModel(); +resultingModel.addModel(filterModel); +resultingModel.addModel(upscalerModel); +resultingModel.compile(); - const count = from.length; - const inputData = new Array(count); - for (let k = 0; k < count; k++) { - inputData[k] = encFrom.map( - (arr, i) => arr.map((value, j) => - value + (encTo[i][j] - value) * k / (count - 1) - ) - ); - } +const gFilter = new GpuModelWrapper(filterModel, {batchSize}); +const gUpscaler = new GpuModelWrapper(upscalerModel, {batchSize: batchSizeUpscaler}); +const gResult = new GpuModelWrapper(resultingModel, {batchSize}); - return pDecoder.compute(inputData.flat()); -} +let quitRequested = false; +process.on("SIGINT", async () => quitRequested = true); async function saveModel() { const savePath = path.join(outPath, "models"); - await ModelUtils.saveModels({autoencoder: chain}, savePath); + await ModelUtils.saveModels({autoencoder: resultingModel}, savePath, true); console.log("Generate final sample set..."); - const genFrom = Matrix.fill(() => dataSet[Math.floor(Math.random() * dataSet.length)], finalSampleSize); - const genTo = Matrix.fill(() => dataSet[Math.floor(Math.random() * dataSet.length)], finalSampleSize); + const finalGenerated = await ImageUtils.processMultiChannelDataParallel( + gResult, testColorfulDataSet.slice(0, finalSampleSize ** 2), testImageChannel + ); - const generated = await _decode(genFrom, genTo); - await ModelUtils.saveGeneratedModelsSamples("autoencoder", savePath, generated, - {channel: imageChannel, count: finalSampleSize, scale: sampleScale}); + await ModelUtils.saveGeneratedModelsSamples("autoencoder", savePath, finalGenerated, + {channel: testImageChannel, scale: 2, count: finalSampleSize}); } -let quitRequested = false; -process.on("SIGINT", async () => quitRequested = true); - -const genFrom = Matrix.fill(() => dataSet[Math.floor(Math.random() * dataSet.length)], epochSampleSize); -const genTo = Matrix.fill(() => dataSet[Math.floor(Math.random() * dataSet.length)], epochSampleSize); - console.log("Training..."); -for (const epoch of ProgressUtils.progress(trainIterations)) { - await pChain.train(dataSet, dataSet, {epochs: epochsPerIter}); +const epochTestData = testColorfulDataSet.slice(0, epochSampleSize ** 2); +const gsEpochTestData = Matrix.fill(() => smallDataSet[Math.floor(Math.random() * smallDataSet.length)], epochSampleSize ** 2); - const generated = await _decode(genFrom, genTo); - await ModelUtils.saveGeneratedModelsSamples(epoch, outPath, generated, - {channel: imageChannel, count: epochSampleSize, scale: sampleScale, time: false, prefix: "autoencoder"}); +for (const epoch of ProgressUtils.progress(trainIterations)) { + gFilter.train(smallDataSet, smallDataSet, {epochs: epochsPerIterFilter}); + gUpscaler.train(smallDataSet, bigDataSet, {epochs: epochsPerIterUpscaler}); + + const filter = gFilter.compute(gsEpochTestData); + await ModelUtils.saveGeneratedModelsSamples(epoch, outPath, filter, + { + channel: imageChannel, + count: epochSampleSize, + scale: 4, + time: false, + prefix: "autoencoder", + suffix: "filter" + }); + + const upscalerGenerated = gUpscaler.compute(gsEpochTestData); + await ModelUtils.saveGeneratedModelsSamples(epoch, outPath, upscalerGenerated, + { + channel: imageChannel, + count: epochSampleSize, + scale: 2, + time: false, + prefix: "autoencoder", + suffix: "upscaler" + }); + + const finalGenerated = await ImageUtils.processMultiChannelDataParallel(gResult, epochTestData, testImageChannel); + await ModelUtils.saveGeneratedModelsSamples(epoch, outPath, finalGenerated, + { + channel: testImageChannel, + count: epochSampleSize, + scale: 2, + time: false, + prefix: "autoencoder", + suffix: "final", + }); if (quitRequested) break; } await saveModel(); -pChain.destroy(); -pEncoder.destroy(); -pDecoder.destroy(); +gFilter.destroy(); +gUpscaler.destroy(); +gResult.destroy(); diff --git a/examples/src/utils/model.js b/examples/src/utils/model.js index 46872c1..13eb9b9 100644 --- a/examples/src/utils/model.js +++ b/examples/src/utils/model.js @@ -5,7 +5,9 @@ import {JsonStreamStringify} from 'json-stream-stringify'; import { GanSerialization, GenerativeAdversarialModel, - UniversalModelSerializer + UniversalModelSerializer, + BinarySerializer, + TensorType } from "mind-net.js"; import * as Image from "./image.js"; @@ -33,12 +35,18 @@ export async function saveModel(model, fileName) { }); } +export async function saveModelBinary(model, fileName) { + const dump = BinarySerializer.save(model, TensorType.F32); + return await CommonUtils.promisify(fs.writeFile, fileName, dump); +} + /** * @param {{[key: string]: IModel|GenerativeAdversarialModel}} models * @param {string} outPath + * @param {boolean} [binary=false] * @return {Promise} */ -export async function saveModels(models, outPath) { +export async function saveModels(models, outPath, binary = false) { const time = new Date().toISOString(); const count = Object.keys(models).length; console.log(`Saving ${count} models...`); @@ -50,8 +58,12 @@ export async function saveModels(models, outPath) { for (const [key, model] of Object.entries(models)) { const epoch = model.epoch ?? model.ganChain?.epoch ?? 0; - const fileName = path.join(outPath, `${key}_${time}_${epoch}.json`); - await saveModel(model, fileName); + const fileName = path.join(outPath, `${key}_${time}_${epoch})`); + if (binary) { + await saveModelBinary(model, fileName + ".bin"); + } else { + await saveModel(model, fileName + ".json"); + } } }