基于java打造的深度学习框架，帮助你快速搭建神经网络，实现训练或测试模型，支持多GPU训练

package com.omega.engine.optimizer; import com.omega.common.data.Tensor; import com.omega.common.data.utils.DataTransforms; import com.omega.common.utils.*; import com.omega.engine.check.BaseCheck; import com.omega.engine.gpu.CUDAModules; import com.omega.engine.nn.data.BaseData; import com.omega.engine.nn.grad.GradClipping; import com.omega.engine.nn.layer.Layer; import com.omega.engine.nn.network.*; import com.omega.engine.nn.network.vae.*; import com.omega.engine.nn.network.vqgan.LPIPS; import com.omega.engine.nn.network.vqgan.PatchGANDiscriminator; import com.omega.engine.optimizer.lr.LearnRateUpdate; import com.omega.example.diffusion.utils.DiffusionImageDataLoader; import com.omega.example.rnn.data.OneHotDataLoader; import com.omega.example.rnn.data.RNNDataLoader; import com.omega.example.sd.utils.SDImageDataLoader; import com.omega.example.sd.utils.SDImageDataLoaderEN; import com.omega.example.transformer.utils.ModelUtils; import com.omega.example.yolo.data.BaseDataLoader; import com.omega.example.yolo.data.DetectionDataLoader; import com.omega.example.yolo.utils.YoloLabelUtils; import jcuda.driver.JCudaDriver; import jcuda.runtime.JCuda; import java.math.BigDecimal; import java.util.Arrays; import java.util.HashMap; import java.util.Map; /** * Mini Batch Stochastic Gradient Descent * * @author Administrator */ public class MBSGDOptimizer extends Optimizer { private YoloLabelUtils u; public MBSGDOptimizer(Network network, int trainTime, float error, int batchSize, boolean warmUp) throws Exception { super(network, batchSize, trainTime, error, warmUp); // TODO Auto-generated constructor stub this.batchSize = batchSize; this.loss = new Tensor(batchSize, this.network.oChannel, this.network.oHeight, this.network.oWidth); this.lossDiff = new Tensor(batchSize, this.network.oChannel, this.network.oHeight, this.network.oWidth); } public MBSGDOptimizer(String sid, Network network, int trainTime, float error, int batchSize, boolean warmUp) throws Exception { super(network, batchSize, trainTime, error, warmUp); // TODO Auto-generated constructor stub this.setSid(sid); this.batchSize = batchSize; this.loss = new Tensor(batchSize, this.network.oChannel, this.network.oHeight, this.network.oWidth); this.lossDiff = new Tensor(batchSize, this.network.oChannel, this.network.oHeight, this.network.oWidth); } public MBSGDOptimizer(Network network, int trainTime, float error, int batchSize, LearnRateUpdate learnRateUpdate, boolean warmUp) throws Exception { super(network, batchSize, trainTime, error, warmUp); // TODO Auto-generated constructor stub this.batchSize = batchSize; this.loss = new Tensor(batchSize, this.network.oChannel, this.network.oHeight, this.network.oWidth); this.lossDiff = new Tensor(batchSize, this.network.oChannel, this.network.oHeight, this.network.oWidth); this.learnRateUpdate = learnRateUpdate; } public MBSGDOptimizer(Network network, int trainTime, float error, int batchSize, LearnRateUpdate learnRateUpdate, boolean warmUp, BaseCheck check) throws Exception { super(network, batchSize, trainTime, error, warmUp); // TODO Auto-generated constructor stub this.batchSize = batchSize; this.loss = new Tensor(batchSize, this.network.oChannel, this.network.oHeight, this.network.oWidth); this.lossDiff = new Tensor(batchSize, this.network.oChannel, this.network.oHeight, this.network.oWidth); this.learnRateUpdate = learnRateUpdate; this.check = check; } public MBSGDOptimizer(String sid, Network network, int trainTime, float error, int batchSize, LearnRateUpdate learnRateUpdate, boolean warmUp) throws Exception { super(network, batchSize, trainTime, error, warmUp); // TODO Auto-generated constructor stub this.setSid(sid); this.batchSize = batchSize; this.loss = new Tensor(batchSize, this.network.oChannel, this.network.oHeight, this.network.oWidth); this.lossDiff = new Tensor(batchSize, this.network.oChannel, this.network.oHeight, this.network.oWidth); this.learnRateUpdate = learnRateUpdate; } public static void q_mean_variance(Tensor x_0, Tensor x_t, Tensor t, float[] posterior_mean_coef1, float[] posterior_mean_coef2, float[] posterior_mean) { for (int b = 0; b < x_t.number; b++) { for (int i = 0; i < x_t.getOnceSize(); i++) { int idx = b * x_t.getOnceSize() + i; posterior_mean[idx] = posterior_mean_coef1[b] * x_0.data[idx] - posterior_mean_coef2[b] * x_t.data[idx]; } } } public static void sample_prev_timestep(DiffusionUNetCond2 network, Tensor condInput, Tensor xt, Tensor t, Tensor x0, int timestep, float[] a, float[] b, float[] betas, float[] alphas, float[] alphas_bar) { for (int i = 0; i < xt.number; i++) { t.data[i] = timestep; } t.hostToDevice(); Tensor noisePred = null; if (condInput != null) { noisePred = network.forward(xt, t, condInput); } else { noisePred = network.forward(xt, t); } noisePred.syncHost(); JCuda.cudaDeviceSynchronize(); if (timestep > 0) { System.err.println("timestep:" + timestep); float var = (1.0f - alphas_bar[timestep - 1]) / (1.0f - alphas_bar[timestep]) * betas[timestep]; float sigma = (float) Math.pow(var, 0.5); float[] noise = RandomUtils.gaussianRandom(noisePred.dataLength, 1.0f); for (int i = 0; i < xt.dataLength; i++) { xt.data[i] = (float) ((xt.data[i] - (betas[timestep] * noisePred.data[i]) / b[timestep]) / Math.sqrt(alphas[timestep])) + sigma * noise[i]; } } else { /** * mean */ for (int i = 0; i < xt.dataLength; i++) { xt.data[i] = (float) ((xt.data[i] - (betas[timestep] * noisePred.data[i]) / b[timestep]) / Math.sqrt(alphas[timestep])); } } xt.hostToDevice(); } public static void showImgs(String outputPath, Tensor input) { ImageUtils utils = new ImageUtils(); if (input.isHasGPU()) { input.syncHost(); } for (int b = 0; b < input.number; b++) { float[] once = input.getByNumber(b); // once = MatrixOperation.add(once, 0.5f); utils.createRGBImage(outputPath + b + ".png", "png", ImageUtils.color2rgb2(once, input.channel, input.height, input.width, true), input.height, input.width, null, null); } } public static void showImgs(String outputPath, Tensor input, String it, float[] mean, float[] std) { ImageUtils utils = new ImageUtils(); for (int b = 0; b < input.number; b++) { float[] once = input.getByNumber(b); utils.createRGBImage(outputPath + it + "_" + b + ".png", "png", ImageUtils.color2rgb2(once, input.channel, input.height, input.width, true, mean, std), input.height, input.width, null, null); } } public static void showImgs(String outputPath, Tensor input, String it, float[] mean, float[] std, String[] labels) { ImageUtils utils = new ImageUtils(); if (labels != null) { for (int b = 0; b < input.number; b++) { float[] once = input.getByNumber(b); String title = labels[b]; if (title.length() > 30) { title = title.substring(0, 30); } utils.createRGBImage(outputPath + it + "_[" + title + "]" + b + ".png", "png", ImageUtils.color2rgb2(once, input.channel, input.height, input.width, true, mean, std), input.height, input.width, null, null); } } else { for (int b = 0; b < input.number; b++) { float[] o