All Downloads are FREE. Search and download functionalities are using the official Maven repository.

com.simiacryptus.mindseye.layers.cudnn.ImgLinearSubnetLayer Maven / Gradle / Ivy

/*
 * Copyright (c) 2019 by Andrew Charneski.
 *
 * The author licenses this file to you under the
 * Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance
 * with the License.  You may obtain a copy
 * of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

package com.simiacryptus.mindseye.layers.cudnn;

import com.google.gson.JsonArray;
import com.google.gson.JsonObject;
import com.simiacryptus.lang.ref.ReferenceCounting;
import com.simiacryptus.lang.ref.ReferenceCountingBase;
import com.simiacryptus.mindseye.lang.*;
import com.simiacryptus.mindseye.lang.cudnn.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.IntStream;
import java.util.stream.Stream;

@SuppressWarnings("serial")
public class ImgLinearSubnetLayer extends LayerBase implements MultiPrecision {

  private static final Logger logger = LoggerFactory.getLogger(ImgLinearSubnetLayer.class);
  private final List legs = new ArrayList<>();
  private Precision precision = CudaSettings.INSTANCE().defaultPrecision;
  private boolean parallel = true;

  public ImgLinearSubnetLayer() {
    super();
  }

  protected ImgLinearSubnetLayer(@Nonnull final JsonObject json, Map rs) {
    super(json);
    this.precision = Precision.valueOf(json.getAsJsonPrimitive("precision").getAsString());
    setParallel(json.get("parallel").getAsBoolean());
    JsonArray jsonArray = json.get("legs").getAsJsonArray();
    for (int i = 0; i < jsonArray.size(); i++) {
      legs.add(new SubnetLeg(jsonArray.get(i).getAsJsonObject(), rs));
    }
  }

  public static ImgLinearSubnetLayer fromJson(@Nonnull final JsonObject json, Map rs) {
    return new ImgLinearSubnetLayer(json, rs);
  }

  public List getLegs() {
    return legs;
  }

  public ImgLinearSubnetLayer add(final int from, final int to, final Layer layer) {
    getLegs().add(new SubnetLeg(layer, from, to));
    return this;
  }

  @Override
  protected void _free() {
    super._free();
    legs.stream().forEach(ReferenceCounting::freeRef);
  }

  @Nullable
  @Override
  public Result evalAndFree(@Nonnull final Result... inObj) {
    assert 1 == inObj.length;
    Result input = inObj[0];
    TensorList inputData = input.getData();
    @Nonnull final int[] inputDims = inputData.getDimensions();
    assert 3 == inputDims.length;
    int length = inputData.length();
    int maxBand = legs.stream().mapToInt(x -> x.toBand).max().getAsInt();
    assert maxBand == inputDims[2] : maxBand + " != " + inputDims[2];
    assert IntStream.range(0, maxBand).allMatch(i ->
        1 == legs.stream().filter(x -> x.fromBand <= i && x.toBand > i).count()
    );
    CudaTensor passback = CudaSystem.run(gpu -> {
      return CudaTensor.wrap(
          gpu.allocate(inputData.getElements() * precision.size, MemoryType.Device, true),
          gpu.newTensorDescriptor(precision, length, inputDims[2], inputDims[1], inputDims[0]), precision);
    });
    inputData.freeRef();
    AtomicInteger counter = new AtomicInteger(0);
    Result[] legResults;
    try {
      legResults = legs.stream().map(leg -> {
        ImgBandSelectLayer imgBandSelectLayer = new ImgBandSelectLayer(leg.fromBand, leg.toBand);
        input.addRef();
        TensorList legData = imgBandSelectLayer.eval(input).getDataAndFree();
        imgBandSelectLayer.freeRef();
        passback.addRef();
        return leg.inner.evalAndFree(new Result(legData, (DeltaSet ctx, TensorList delta) -> {
          int[] outputDimensions = delta.getDimensions();
          int[] inputDimensions = inputDims;
          synchronized (passback) {
            CudaSystem.run(gpu -> {
              @Nonnull final CudaDevice.CudaTensorDescriptor viewDescriptor = gpu.newTensorDescriptor(
                  precision, length, outputDimensions[2], outputDimensions[1], outputDimensions[0], //
                  inputDimensions[2] * inputDimensions[1] * inputDimensions[0], //
                  inputDimensions[1] * inputDimensions[0], //
                  inputDimensions[0], //
                  1);
              final int byteOffset = viewDescriptor.cStride * leg.fromBand * precision.size;
              assert delta.length() == length;
              assert passback.getDeviceId() == gpu.getDeviceId();
              //assert error.stream().flatMapToDouble(x-> Arrays.stream(x.getData())).allMatch(Double::isFinite);
              @Nullable final CudaTensor deltaTensor = gpu.getTensor(delta, precision, MemoryType.Device, true);
              @Nonnull final CudaMemory passbackBuffer = passback.getMemory(gpu);
              CudaMemory errorPtrMemory = deltaTensor.getMemory(gpu);
              passbackBuffer.synchronize();
              gpu.cudnnTransformTensor(
                  precision.getPointer(1.0), deltaTensor.descriptor.getPtr(), errorPtrMemory.getPtr(),
                  precision.getPointer(0.0), viewDescriptor.getPtr(), passbackBuffer.getPtr().withByteOffset(byteOffset)
              );
              errorPtrMemory.dirty();
              passbackBuffer.dirty();
              Stream.of(deltaTensor, viewDescriptor, passbackBuffer, errorPtrMemory).forEach(ReferenceCounting::freeRef);
            }, passback);
          }
          if (counter.incrementAndGet() >= legs.size()) {
            counter.set(0);
            input.accumulate(ctx, CudaTensorList.create(passback, length, inputDims, precision));
          }
          delta.freeRef();
        }) {
          @Override
          protected void _free() {
            input.freeRef();
            passback.freeRef();
            super._free();
          }
        });
      }).toArray(i -> new Result[i]);
    } finally {
      input.freeRef();
      passback.freeRef();
    }
    SumInputsLayer sumInputsLayer = new SumInputsLayer().setParallel(parallel).setPrecision(precision);
    try {
      return sumInputsLayer.evalAndFree(legResults);
    } finally {
      sumInputsLayer.freeRef();
    }
  }

  @Nonnull
  @Override
  public JsonObject getJson(Map resources, DataSerializer dataSerializer) {
    @Nonnull final JsonObject json = super.getJsonStub();
    json.addProperty("precision", precision.name());
    json.addProperty("parallel", isParallel());
    JsonArray jsonArray = new JsonArray();
    legs.stream().map(x -> x.getJson(resources, dataSerializer)).forEach(jsonArray::add);
    json.add("legs", jsonArray);
    return json;
  }

  @Nonnull
  @Override
  public List state() {
    return new ArrayList<>();
  }

  @Override
  public Precision getPrecision() {
    return precision;
  }

  @Nonnull
  @Override
  public ImgLinearSubnetLayer setPrecision(Precision precision) {
    this.precision = precision;
    return this;
  }

  @Nonnull
  @Override
  public Layer setFrozen(final boolean frozen) {
    legs.stream().map(x -> x.inner).forEach(x -> x.setFrozen(frozen));
    return super.setFrozen(frozen);
  }

  public boolean isParallel() {
    return parallel;
  }

  public ImgLinearSubnetLayer setParallel(boolean parallel) {
    this.parallel = parallel;
    return this;
  }

  public static class SubnetLeg extends ReferenceCountingBase {

    private final Layer inner;
    private final int fromBand;
    private final int toBand;

    public SubnetLeg(final Layer inner, final int fromBand, final int toBand) {
      this.inner = inner;
      this.fromBand = fromBand;
      this.toBand = toBand;
      this.inner.addRef();
    }

    protected SubnetLeg(@Nonnull final JsonObject json, Map rs) {
      fromBand = json.getAsJsonPrimitive("fromBand").getAsInt();
      toBand = json.getAsJsonPrimitive("toBand").getAsInt();
      inner = Layer.fromJson(json.getAsJsonObject("network"), rs);
    }

    @Override
    protected void _free() {
      super._free();
      inner.freeRef();
    }

    @Nonnull
    public JsonObject getJson(Map resources, DataSerializer dataSerializer) {
      @Nonnull final JsonObject json = new JsonObject();
      json.addProperty("fromBand", fromBand);
      json.addProperty("toBand", toBand);
      json.add("network", inner.getJson(resources, dataSerializer));
      return json;
    }

  }
}




© 2015 - 2025 Weber Informatics LLC | Privacy Policy