[자바스크립트] Node.js - Tensorflow.js Generator를 이용한 데이터셋 만들기

파이썬의 케라스처럼 Tensorflow.js에도 data generator가 있습니다.

 

Javascript에서는 function* 처럼 표현을 하는데요, 사용하면서 주의점이 하나 있는데

 

일반 함수처럼 인자를 받으면 제대로 작동하지 않습니다.

 -> function* imageGenerator(x)  실행 안됨...

 

따라서 imgPath 라는 이미지파일 경로를 저장한 배열이 전역변수로 필요합니다.

 

let imgPath = [];
async function* imageGenerator() {
    for (let i = 0; i < imgPath.length; i++) {
        console.log(i);
        const img = await loadImage(imgPath[i]);
        const tens = await tf.browser.fromPixels(img)
            .resizeNearestNeighbor([imgSize, imgSize])
            .div(255.0);
        yield tens;
    }
}

async function* labelGenerator() {
    const labels = imgPath.map(path => {
        const labelSplit = path.split("\\");
        const label = labelSplit[labelSplit.length - 2];
        return label;
    });
    const uniq = labels.reduce(function (a, b) {
        if (a.indexOf(b) < 0) a.push(b);
        return a;
    }, []);

    for (let j = 0; j < labels.length; j++) {
        for (let i = 0; i < uniq.length; i++) {
            if(labels[j] == uniq[i]){
                const lbl = await tf.oneHot(i, uniq.length);
                yield lbl;
            }
        }
    }
}

위의 함수를 이용해 최종 데이터셋을 만드는 코드는 아래와 같습니다.

async function run() {
    const xs = await tf.data.generator(imageGenerator);
    const ys = await tf.data.generator(labelGenerator);
    const ds = await tf.data.zip({xs,ys}).batch(10).shuffle(imgPath.length);
}

 

데이터셋을 만들고, batch를 10으로하고 전체에 대해 shuffle을 수행했습니다.

model.fit의 경우 데이터셋을 default로 shuffle을 적용해 주지만

데이터셋을 이용한 경우 꼭 shuffle을 해줘야 합니다. 이 부분을 몰랐어서 성능이 안나왔었네요 ㅎㅎ;

 

전체 코드는 다음 링크를 참고해 주시기 바랍니다.

2021/01/23 - [자바스크립트] - Node.js - Tensorflow.js를 이용한 이미지 분류 예제

 

* 이 글은 티스토리 카카오계정 연동정책으로 인해 이전 블로그(오코취) 글을 옮겨왔습니다.

[끝].