Java AffinityManager类代码示例

import org.nd4j.linalg.api.concurrency.AffinityManager; //导入依赖的package包/类
@Override
public INDArray compress(INDArray array) {
    //logger.info("Threshold [{}] compression", threshold);

    Nd4j.getExecutioner().commit();
    Nd4j.getAffinityManager().ensureLocation(array, AffinityManager.Location.HOST);

    DataBuffer buffer = compress(array.data());
    if (buffer == null)
        return null;

    INDArray dup = Nd4j.createArrayFromShapeBuffer(buffer, array.shapeInfoDataBuffer());
    dup.markAsCompressed(true);

    return dup;
}
 

import org.nd4j.linalg.api.concurrency.AffinityManager; //导入依赖的package包/类
/**
 * Setup the given byte buffer
 * for serialization (note that this is for uncompressed INDArrays)
 * 4 bytes int for rank
 * 4 bytes for data opType
 * shape buffer
 * data buffer
 *
 * @param arr the array to setup
 * @param allocated the byte buffer to setup
 * @param rewind whether to rewind the byte buffer or nt
 */
public static void doByteBufferPutUnCompressed(INDArray arr, ByteBuffer allocated, boolean rewind) {
    // ensure we send data to host memory
    Nd4j.getExecutioner().commit();
    Nd4j.getAffinityManager().ensureLocation(arr, AffinityManager.Location.HOST);

    ByteBuffer buffer = arr.data().pointer().asByteBuffer().order(ByteOrder.nativeOrder());
    ByteBuffer shapeBuffer = arr.shapeInfoDataBuffer().pointer().asByteBuffer().order(ByteOrder.nativeOrder());
    //2 four byte ints at the beginning
    allocated.putInt(arr.rank());
    //put data opType next so its self describing
    allocated.putInt(arr.data().dataType().ordinal());
    allocated.put(shapeBuffer);
    allocated.put(buffer);
    if (rewind)
        allocated.rewind();
}
 

import org.nd4j.linalg.api.concurrency.AffinityManager; //导入依赖的package包/类
@Override
public void call(DataSet dataSet) {
    if (dataSet != null) {
        if (dataSet.getFeatures() != null)
            Nd4j.getAffinityManager().ensureLocation(dataSet.getFeatures(), AffinityManager.Location.DEVICE);

        if (dataSet.getLabels() != null)
            Nd4j.getAffinityManager().ensureLocation(dataSet.getLabels(), AffinityManager.Location.DEVICE);

        if (dataSet.getFeaturesMaskArray() != null)
            Nd4j.getAffinityManager().ensureLocation(dataSet.getFeaturesMaskArray(),
                            AffinityManager.Location.DEVICE);

        if (dataSet.getLabelsMaskArray() != null)
            Nd4j.getAffinityManager().ensureLocation(dataSet.getLabelsMaskArray(), AffinityManager.Location.DEVICE);
    }
}
 

import org.nd4j.linalg.api.concurrency.AffinityManager; //导入依赖的package包/类
@Override
public List<Writable> next() {
    if (iter != null) {
        List<Writable> ret;
        File image = iter.next();
        currentFile = image;

        if (image.isDirectory())
            return next();
        try {
            invokeListeners(image);
            INDArray row = imageLoader.asMatrix(image);
            Nd4j.getAffinityManager().ensureLocation(row, AffinityManager.Location.DEVICE);
            ret = RecordConverter.toRecord(row);
            if (appendLabel || writeLabel){
                if( labelGenerator.inferLabelClasses()){
                    //Standard classification use case (i.e., handle String -> integer conversion
                    ret.add(new IntWritable(labels.indexOf(getLabel(image.getPath()))));
                } else {
                    //Regression use cases, and PathLabelGenerator instances that already map to integers
                    ret.add(labelGenerator.getLabelForPath(image.getPath()));
                }
            }
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
        return ret;
    } else if (record != null) {
        hitImage = true;
        invokeListeners(record);
        return record;
    }
    throw new IllegalStateException("No more elements");
}
 

import org.nd4j.linalg.api.concurrency.AffinityManager; //导入依赖的package包/类
@Override
public DataBuffer compress(DataBuffer buffer) {
    INDArray temp = Nd4j.createArrayFromShapeBuffer(buffer, Nd4j.getShapeInfoProvider().createShapeInformation(new int[]{1, (int) buffer.length()}));
    double max = temp.amaxNumber().doubleValue();

    int cntAbs = temp.scan(Conditions.absGreaterThanOrEqual(max - (max * threshold))).intValue();

    long originalLength = buffer.length() * Nd4j.sizeOfDataType(buffer.dataType());
    int compressedLength = cntAbs + 3;
    // first 3 elements contain header
    IntPointer pointer = new IntPointer(compressedLength);
    pointer.put(0, cntAbs);
    pointer.put(1, (int) buffer.length());
    pointer.put(2, Float.floatToIntBits(threshold)); // please note, this value will be ovewritten anyway

    CompressionDescriptor descriptor = new CompressionDescriptor();
    descriptor.setCompressedLength(compressedLength * 4); // sizeOf(INT)
    descriptor.setOriginalLength(originalLength);
    descriptor.setOriginalElementSize(Nd4j.sizeOfDataType(buffer.dataType()));
    descriptor.setNumberOfElements(buffer.length());

    descriptor.setCompressionAlgorithm(getDescriptor());
    descriptor.setCompressionType(getCompressionType());

    CompressedDataBuffer cbuff = new CompressedDataBuffer(pointer, descriptor);

    Nd4j.getNDArrayFactory().convertDataEx(getBufferTypeEx(buffer), buffer.addressPointer(), DataBuffer.TypeEx.FTHRESHOLD, pointer, buffer.length());

    Nd4j.getAffinityManager().tagLocation(buffer, AffinityManager.Location.HOST);

    return cbuff;
}
 

import org.nd4j.linalg.api.concurrency.AffinityManager; //导入依赖的package包/类
@Override
public DataBuffer compress(DataBuffer buffer) {
    INDArray temp = Nd4j.createArrayFromShapeBuffer(buffer, Nd4j.getShapeInfoProvider().createShapeInformation(new int[]{1, (int) buffer.length()}).getFirst());
    double max = temp.amaxNumber().doubleValue();

    int cntAbs = temp.scan(Conditions.absGreaterThanOrEqual(max - (max * threshold))).intValue();

    long originalLength = buffer.length() * Nd4j.sizeOfDataType(buffer.dataType());
    int compressedLength = cntAbs + 4;
    // first 3 elements contain header
    IntPointer pointer = new IntPointer(compressedLength);
    pointer.put(0, cntAbs);
    pointer.put(1, (int) buffer.length());
    pointer.put(2, Float.floatToIntBits(threshold)); // please note, this value will be ovewritten anyway
    pointer.put(3, 0);

    CompressionDescriptor descriptor = new CompressionDescriptor();
    descriptor.setCompressedLength(compressedLength * 4); // sizeOf(INT)
    descriptor.setOriginalLength(originalLength);
    descriptor.setOriginalElementSize(Nd4j.sizeOfDataType(buffer.dataType()));
    descriptor.setNumberOfElements(buffer.length());

    descriptor.setCompressionAlgorithm(getDescriptor());
    descriptor.setCompressionType(getCompressionType());

    CompressedDataBuffer cbuff = new CompressedDataBuffer(pointer, descriptor);

    Nd4j.getNDArrayFactory().convertDataEx(getBufferTypeEx(buffer), buffer.addressPointer(), DataBuffer.TypeEx.FTHRESHOLD, pointer, buffer.length());

    Nd4j.getAffinityManager().tagLocation(buffer, AffinityManager.Location.HOST);

    return cbuff;
}
 

import org.nd4j.linalg.api.concurrency.AffinityManager; //导入依赖的package包/类
@Override
public List<Writable> next(int num) {
    List<File> files = new ArrayList<>(num);
    List<List<ImageObject>> objects = new ArrayList<>(num);

    for (int i = 0; i < num && hasNext(); i++) {
        File f = iter.next();
        this.currentFile = f;
        if (!f.isDirectory()) {
            files.add(f);
            objects.add(labelProvider.getImageObjectsForPath(f.getPath()));
        }
    }


    int nClasses = labels.size();

    INDArray outImg = Nd4j.create(files.size(), channels, height, width);
    INDArray outLabel = Nd4j.create(files.size(), 4 + nClasses, gridH, gridW);

    int exampleNum = 0;
    for (int i = 0; i < files.size(); i++) {
        File imageFile = files.get(i);
        this.currentFile = imageFile;
        try {
            this.invokeListeners(imageFile);
            Image image = this.imageLoader.asImageMatrix(imageFile);
            this.currentImage = image;
            Nd4j.getAffinityManager().ensureLocation(image.getImage(), AffinityManager.Location.DEVICE);

            outImg.put(new INDArrayIndex[]{point(exampleNum), all(), all(), all()}, image.getImage());

            List<ImageObject> objectsThisImg = objects.get(exampleNum);

            label(image, objectsThisImg, outLabel, exampleNum);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }

        exampleNum++;
    }

    return Arrays.<Writable>asList(new NDArrayWritable(outImg), new NDArrayWritable(outLabel));
}
 

import org.nd4j.linalg.api.concurrency.AffinityManager; //导入依赖的package包/类
@Override
public INDArray thresholdEncode(INDArray input, double threshold, Integer boundary) {

    MatchCondition condition = new MatchCondition(input, Conditions.absGreaterThanOrEqual(threshold));
    int cntAbs = Nd4j.getExecutioner().exec(condition, Integer.MAX_VALUE).getInt(0);

    if (cntAbs < 2)
        return null;

    if (boundary != null)
        cntAbs = Math.min(cntAbs, boundary);

    DataBuffer buffer = input.data();

    long originalLength = buffer.length() * Nd4j.sizeOfDataType(buffer.dataType());
    int compressedLength = cntAbs + 4;
    // first 3 elements contain header

    DataBuffer encodedBuffer = Nd4j.getMemoryManager().getCurrentWorkspace() == null ? Nd4j.getDataBufferFactory().createInt(4+cntAbs, false) : Nd4j.getDataBufferFactory().createInt(4+cntAbs, false, Nd4j.getMemoryManager().getCurrentWorkspace());

    encodedBuffer.put(0, cntAbs);
    encodedBuffer.put(1, (int) buffer.length());
    encodedBuffer.put(2, Float.floatToIntBits((float) threshold));

    // format id
    encodedBuffer.put(3, ThresholdCompression.FLEXIBLE_ENCODING);

    CompressionDescriptor descriptor = new CompressionDescriptor();
    descriptor.setCompressedLength(compressedLength * 4); // sizeOf(INT)
    descriptor.setOriginalLength(originalLength);
    descriptor.setOriginalElementSize(Nd4j.sizeOfDataType(buffer.dataType()));
    descriptor.setNumberOfElements(buffer.length());

    descriptor.setCompressionAlgorithm("THRESHOLD");
    descriptor.setCompressionType(CompressionType.LOSSLESS);

    //CompressedDataBuffer cbuff = new CompressedDataBuffer(pointer, descriptor);

    Nd4j.getNDArrayFactory().convertDataEx(AbstractCompressor.getBufferTypeEx(buffer), buffer.addressPointer(), DataBuffer.TypeEx.THRESHOLD, encodedBuffer.addressPointer(), buffer.length());

    Nd4j.getAffinityManager().tagLocation(buffer, AffinityManager.Location.HOST);

    return Nd4j.createArrayFromShapeBuffer(encodedBuffer, input.shapeInfoDataBuffer());
}
 

import org.nd4j.linalg.api.concurrency.AffinityManager; //导入依赖的package包/类
/**
 *
 * @return
 */
public static AffinityManager getAffinityManager() {
    return affinityManager;
}