Начало работы с нейронным шейдингом для игр

В последние 25 лет рендеринг в реальном времени развивался благодаря постоянным улучшениям аппаратного обеспечения. Основная задача всегда заключалась в создании изображения максимальной детализации за 16 миллисекунд. Это стимулировало инновации в графическом оборудовании, конвейерах и системах рендеринга.

Однако замедление закона Мура требует новых вычислительных архитектур, чтобы справляться с растущими требованиями приложений реального времени. Аналогично, когда классические методы графики достигают пределов, нужны свежие подходы для повышения качества изображения и скорости. Возникает ключевая задача: как дальше совершенствовать рендеринг в реальном времени, не полагаясь только на традиционные улучшения железа?

Нейронный шейдинг предлагает свежий путь — внедрение обучаемых моделей прямо в графический конвейер для достижения выдающегося качества и производительности. Эта технология использует специализированное AI-оборудование, например Tensor Cores от NVIDIA, чтобы запускать нейронные сети эффективно в реальном времени.

В этой статье разберем основы и покажем, как войти в эту меняющую всё технологию.

Что такое нейронный шейдинг?

В основе нейронного шейдинга лежит идея сделать часть графического конвейера обучаемой. Это может касаться любых параметров, которые настраиваются с помощью методов машинного обучения, но особенно перспективны компактные нейронные сети, интегрированные в шейдеры и взаимодействующие с остальной системой рендеринга.

Такие сети работают очень быстро в реальном времени, особенно с ускорением от технологий вроде кооперативных векторов. С одной стороны, это выжимает максимум из текущего оборудования, позволяя отображать больше деталей без уменьшения транзисторов. С другой — обучаемые шейдеры решают задачи, сложные для традиционных методов. Это добавляет в арсенал графики рабочий инструмент, доступный уже сейчас.

Как это меняет подход?

Классическая инженерия подразумевает понимание проблемы, ее решение, кодирование и выполнение. Но некоторые задачи не имеют готовых решений, или они слишком затратны для реального времени, особенно в нейронном шейдинге. Здесь на помощь приходит оптимизация. Вместо прямого решения мы берем известные входы и выходы, обучаем математическую модель, подстраивая параметры шаг за шагом, пока не получим полезное приближение.

Современный нейронный шейдинг использует мощные инструменты вроде Slang — языка шейдинга, который становится ключевым в разработке игр. Его поддерживает Khronos, организация, отвечающая за API вроде OpenGL и Vulkan. Slang обеспечивает совместимость с платформами, целевыми для HLSL, SPIR-V, Metal и других. Он включает современные конструкции вроде generics и, что важно для нейронного шейдинга, автоматическое дифференцирование (autodiff), упрощающее сложные вычисления.

SlangPy — это интерфейс Python для Slang. Он дает доступ к базовым элементам графики, таким как буферы вычислений и текстуры. Полностью кросс-платформенный (под D3D 12, Vulkan, CUDA, Metal), SlangPy предлагает функциональный API для прямого вызова шейдеров Slang из Python.

Для практики посмотрите лабораторию по введению в Slang от SIGGRAPH и материалы для скачивания. Также полезна сессия Birds of a Feather по Slang с обсуждениями сообщества.

С чего начать: пример с mipmap

Возьмем конкретный случай, чтобы проиллюстрировать идеи: генерацию mipmap. Обычные mipmap хорошо справляются с цветными текстурами вроде альбедо, которые сглаживаются простыми фильтрами. Но карты геометрии или топологии сжимаются плохо — нельзя просто усреднить пик с впадиной, чтобы получить плоскую поверхность.

Простой подход приводит к артефактам с шумными бликами и выдуманными поверхностями. Эта проблема изучена давно, есть аналитические решения, например метод Toksvig: он фильтрует карты нормалей, корректируя шероховатость по дисперсии нормалей в уровнях mipmap, учитывая геометрическую сложность на разных масштабах.

Хотя аналитические методы эффективны для узких случаев, они требуют знаний домена и тонкой настройки. Нейронная оптимизация дает универсальный вариант — генерируем mipmap, минимизируя разницу между сжатым рендерингом и эталонным, обучая оптимальные представления без явных формул.

Как работает оптимизация

Процесс оптимизации делится на два этапа:

1. Прямой этап: рендерим идеальный выход традиционно, генерируем результат и измеряем разницу между ними.

2. Обратный этап: вычисляем, как скорректировать входы, чтобы уменьшить ошибку, с помощью автоматического дифференцирования.

Суть в том, что возможности autodiff в Slang автоматически генерируют обратные производные для всего рендеринга на этапе компиляции. Это быстрее и удобнее ручного дифференцирования, и производные всегда синхронизированы с изменениями в прямом коде.

Вот простой пример, как разработчик может подойти к этому в Slang. Это иллюстрация, которую нужно адаптировать под случай: добавить данные текстуры и реализацию BRDF.

// Определяем обучаемые параметры mipmap для материала
struct MaterialParameters {
    GradOutTensor<float3, 2> albedo;
    GradOutTensor<float3, 2> normal;
};

// Дифференцируемая функция рендеринга
[Differentiable]
float3 render(int2 pixel, MaterialParameters material, no_diff float3 light_dir, no_diff float3 view_dir) {
    // Яркий белый свет
    float light_intensity = 5.0;

    // Сэмплируем блестящий BRDF (сегодня шел дождь!)
    float3 brdf_sample = sample_brdf(
        // предполагаем, что BRDF реализован отдельно
        material.get_albedo(pixel), // цвет альбедо
        normalize(light_dir), // направление света
        normalize(view_dir), // направление взгляда
        material.get_normal(pixel), // сэмпл карты нормалей
        0.05, // шероховатость
        0.0, // металлический (не металл)
        1.0 // спекулярный
    );

    // Комбинируем свет с BRDF для цвета пикселя
    return brdf_sample * light_intensity;
}

// Простая функция даунсэмплинга бокс-фильтром
float3 downsample(
    int2 pixel,
    Tensor<float3, 2> source)
{
    float3 res = 0;
    res += source.getv(pixel * 2 + int2(0, 0));
    res += source.getv(pixel * 2 + int2(1, 0));
    res += source.getv(pixel * 2 + int2(0, 1));
    res += source.getv(pixel * 2 + int2(1, 1));
    return res * 0.25;
}

// Функция потерь, сравнивающая mipmap с эталоном
[Differentiable]
float3 loss(
    no_diff int2 pixel,
    no_diff float3 reference,
    MaterialParameters material,
    no_diff float3 light_dir,
    no_diff float3 view_dir)
{
    float3 color = render(pixel, material, light_dir, view_dir);
    float3 error = color - reference;
    return error * error; // Квадратичная ошибка
}

А вот как вызвать этот код через Python/SlangPy:

import slangpy as spy
import pathlib

# Создаем устройство и загружаем модуль Slang
device = spy.create_device(
    include_paths=[
        pathlib.Path(__file__).parent.absolute(),
    ]
)
module = spy.Module.load_from_file(device, "example.slang")

# Загружаем материалы.
albedo_map = spy.Tensor.load_from_image(device, "PavingStones070_2K.diffuse.jpg", linearize=True)
normal_map = spy.Tensor.load_from_image(device, "PavingStones070_2K.normal.jpg", scale=2, offset=-1)

def downsample(source: spy.Tensor, steps: int) -> spy.Tensor:
    for i in range(steps):
        dest = spy.Tensor.empty(
            device=device,
            shape=(source.shape[0] // 2, source.shape[1] // 2),
            dtype=source.dtype)
        module.downsample(spy.call_id(), source, _result=dest)
        source = dest
    return source

# Выделяем тензор для выхода + вызываем функцию рендеринга
output = spy.Tensor.empty_like(albedo_map)
module.render(
    pixel = spy.call_id(),
    material = {
        "albedo": albedo_map,
        "normal": normal_map,
    },
    light_dir = spy.math.normalize(spy.float3(0.2, 0.2, 1.0)),
    view_dir = spy.float3(0, 0, 1),
    _result = output)

# Даунсэмплируем тензор выхода.
output = downsample(output, 2)

# Сохраняем в файл
output_filename = "render_output.png"
output.save_to_image(output_filename)

Пока мы взяли исходную текстуру, сжали ее и сохранили. Теперь нужно вычислить результат из обучаемых параметров и разницу с оригиналом. Цель — обучить меньшие текстуры, чтобы они давали тот же эффект, как полная версия референса. Начнем с расчета потерь:

# Потери между даунсэмплированным выходом полного разрешения (эталон)
# и результатом от входов четвертьного разрешения.
loss_output = spy.Tensor.empty_like(output)
module.loss(
    pixel = spy.call_id(),
    material = {
        "albedo": downsample(albedo_map, 2),
        "normal": downsample(normal_map, 2),
    },
    reference = output,
    light_dir = spy.math.normalize(spy.float3(0.2, 0.2, 1.0)),
    view_dir = spy.float3(0, 0, 1),
    _result = loss_output)

Этот код показывает, насколько результат отличается от желаемого, но не говорит, как подстроить параметры. Для этого нужны градиенты — здесь autodiff Slang поможет. Добавим функцию в Slang:

void calculate_grads(
    int2 pixel,
    MaterialParameters material,
    MaterialParameters ref_material)
{
    float3 light_dir = random_direction(); // Предполагаем реализацию генератора случайных направлений
    float3 view_dir = random_direction();

    // Рендерим высококачественный эталон стандартной функцией
    float3 reference = render(pixel, ref_material, light_dir, view_dir);

    // Обратное распространение
    bwd_diff(loss)(pixel, reference, material, light_dir, view_dir, 1.0);
}

Эта функция использует нашу функцию потерь, чтобы вычислить градиенты для каждого пикселя текстур материала — производную потерь по каждому входу. Градиенты указывают, как обновить текстуры, чтобы уменьшить потери.

Последний шаг — обновить входные текстуры по градиентам и повторить цикл, приближаясь к идеалу. Для примера возьмем простую функцию, но в реальности используют продвинутые оптимизаторы вроде Adam.

void optimizer_step(inout float3 parameter, float3 derivative, float learning_rate) {
    parameter -= learning_rate * derivative;
}

Теперь повторяем шаги в простом цикле Python:

for iteration in range(num_iterations):
    # Шаг 1: Вычисляем градиенты через автоматическое дифференцирование
    module.calculate_grads(
        pixel=spy.call_id(),
        material=trainable_material,
        ref_material=reference_material
    )

    # Шаг 2: Обновляем параметры оптимизатором
    module.optimizer_step(
        pixel=spy.call_id(),
        trainable_material["albedo"],
        learning_rate=learning_rate
    )
    # Повторяем для карты нормалей...

Основы нейронных сетей в шейдерах

Переходя от простой оптимизации параметров, можно встроить целые нейронные сети прямо в шейдеры. Нейронная сеть — это математическая функция, аппроксимирующая сложные связи между входами и выходами. Вместо явного кода для этих связей сеть обучается их извлекать самостоятельно.

Зачем нейронные сети в шейдерах?

Нейронные сети особенно хороши для задач графики:

• Сжатие: Компактная сеть представляет сложные текстуры или материалы с меньшим числом параметров, чем традиционные методы.
• Аппроксимация: Они заменяют дорогие вычисления (сложные модели освещения) быстрыми операциями.
• Генерализация: После обучения сети справляются с вариациями и крайними случаями, трудно программируемыми вручную.
• Оптимизация: Они находят оптимальные решения для задач без аналитических формул или слишком затратных.

Базовые элементы

Основной блок нейронной сети прост: входы (вещественные числа), веса (настраиваемые параметры), смещения (дополнительные настраиваемые) и нелинейная функция активации. Сеть учится, подстраивая веса и смещения, чтобы минимизировать разницу между предсказаниями и желаемыми выходами.

Продолжая с текстурами, создадим сеть, которая принимает UV-координаты и генерирует RGB. С девятью параметрами (шесть весов и три смещения) можно заменить 200 000 чисел традиционной текстуры.

Для этой сети используем гиперболический тангенс (tanh()) как активацию — распространенный выбор. Обучение через шаг оптимизации в фреймворке. Класс NetworkParameters в Python имеет метод optimize(), который вызывает adamOptimize() в модуле Slang. Там реализуется базовый Adam-оптимизатор на GPU.

Вот базовая реализация сети в Slang:

import slangpy;

// Простая функция активации (tanh)
[Differentiable]
float activation(float x) {
    return tanh(x);
}

// Простой Adam-оптимизатор для одного параметра
void adamOptimize(
    inout float primal, // Параметр для оптимизации
    inout float grad, // Градиент
    inout float m_prev, // Первое момент (среднее градиента)
    inout float v_prev, // Второе момент (среднее квадрата градиента)
    float learning_rate, // Скорость обучения
    int iteration) // Номер итерации
{
    const float ADAM_BETA_1 = 0.9;
    const float ADAM_BETA_2 = 0.999;
    const float ADAM_EPSILON = 1e-8;

    // Обновляем моменты
    float m = ADAM_BETA_1 * m_prev + (1.0 - ADAM_BETA_1) * grad;
    float v = ADAM_BETA_2 * v_prev + (1.0 - ADAM_BETA_2) * (grad * grad);
    m_prev = m;
    v_prev = v;

    // Коррекция смещения
    float mHat = m / (1.0f - pow(ADAM_BETA_1, iteration));
    float vHat = v / (1.0f - pow(ADAM_BETA_2, iteration));

    // Обновляем параметр
    primal -= learning_rate * (mHat / (sqrt(vHat) + ADAM_EPSILON));

    // Сбрасываем градиент
    grad = 0;
}

// Параметры сети с поддержкой автоматического дифференцирования
struct NetworkParameters<int Inputs, int Outputs> {
    RWTensor<float, 1> biases;
    RWTensor<float, 2> weights;
    AtomicTensor<float, 1> biases_grad;
    AtomicTensor<float, 2> weights_grad;

    [Differentiable]
    float get_bias(int neuron) {
        return biases.get({neuron});
    }
    [Differentiable]

А вот настройка и обучение сети в Python:

import slangpy as spy
import numpy as np
import pathlib

# Создаем устройство и загружаем модуль Slang
device = spy.create_device(
    include_paths=[
        pathlib.Path(__file__).parent.absolute(),
    ]
)
module = spy.Module.load_from_file(device, "example.slang")

# Обертка Python для структуры NetworkParameters Slang
class NetworkParameters(spy.InstanceList):
    def __init__(self, inputs: int, outputs: int):
        super().__init__(module[f"NetworkParameters<{inputs},{outputs}>"])
        self.inputs = inputs
        self.outputs = outputs

        # Смещения и веса для слоя.
        self.biases = spy.Tensor.from_numpy(device, np.zeros(outputs).astype('float32'))
        self.weights = spy.Tensor.from_numpy(device, np.random.uniform(-0.5, 0.5, (outputs, inputs)).astype('float32'))

        # Градиенты для смещений и весов.
        self.biases_grad = spy.Tensor.zeros_like(self.biases)
        self.weights_grad = spy.Tensor.zeros_like(self.weights)

        # Временные данные для Adam-оптимизатора.
        self.m_biases = spy.Tensor.zeros_like(self.biases)
        self.m_weights = spy.Tensor.zeros_like(self.weights)
        self.v_biases = spy.Tensor.zeros_like(self.biases)
        self.v_weights = spy.Tensor.zeros_like(self.weights)

    # Вызывает функцию 'optimize' Slang для смещений и весов
    def optimize(self, learning_rate: float, optimize_counter: int):
        module.adamOptimize(self.biases, self.biases_grad, self.m_biases, self.v_biases, learning_rate, optimize_counter)
        module.adamOptimize(self.weights, self.weights_grad, self.m_weights, self.v_weights, learning_rate, optimize_counter)

# Создаем параметры сети для слоя с 2 входами и 3 выходами
params = NetworkParameters(2, 3)
print(f"Created NetworkParameters with {params.inputs} inputs and {params.outputs} outputs")
print(f"Biases shape: {params.biases.shape}")
print(f"Weights shape: {params.weights.shape}")
print(f"Initial weights:\n{params.weights.to_numpy()}")

Для сложных сетей легко добавить слои:

// Многослойная сеть для сложной генерации текстур
struct Network {
    NetworkParameters<2, 32> layer0;
    NetworkParameters<32, 32> layer1;
    NetworkParameters<32, 3> layer2;

    [Differentiable]
    float3 eval(no_diff float2 uv) {
        float inputs[2] = {uv.x, uv.y};
        float output0[32] = layer0.forward(inputs);
        [ForceUnroll]
        for (int i = 0; i < 32; ++i)
            output0[i] = activation(output0[i]);

        float output1[32] = layer1.forward(output0);
        [ForceUnroll]
        for (int i = 0; i < 32; ++i)
            output1[i] = activation(output1[i]);

        float output2[3] = layer2.forward(output1);
        [ForceUnroll]
        for (int i = 0; i < 3; ++i)
            output2[i] = activation(output2[i]);

        return float3(output2[0], output2[1], output2[2]);
    }
}

Преимущество в том, что инфраструктура autodiff, работавшая для простой оптимизации, подходит и для обучения сетей. Компилятор генерирует градиенты для всей сети, облегчая обучение моделей генерации текстур.

Ключевые приемы для лучших результатов

Компактные сети требуют тщательной настройки. Лучшие методы зависят от задачи — то, что подходит для текстур, может не сработать для материалов или освещения. Вот приемы, улучшающие результаты для примера с текстурами:

• Функции активации: ReLU — популярная в машинном обучении, пропускает положительные входы и обнуляет отрицательные. Она эффективна и быстра для многих задач нейронного шейдинга, но создает кусочно-линейные выходы из-за порога в ноль, что приводит к треугольным паттернам в 2D-текстурах. Гладкие активации вроде экспоненты дают лучший визуал для генерации текстур. Выбор зависит от случая.

// Альтернативные функции активации
[Differentiable]
float3 smoothActivation(float3 x) {
    return exp(x); // Экспоненциальная для гладкого выхода
}

[Differentiable]
float3 leakyReLU(float3 x) {
    return max(0.1 * x, x); // Leaky ReLU предотвращает мертвые нейроны
}

• Leaky ReLU: Когда ReLU обнуляет отрицательный вход, градиент тоже становится нулевым. При обратном распространении обновления не доходят до весов этого нейрона. Если нейрон часто получает отрицательные входы, он 'умирает' — всегда выдает ноль и не учится. В малых сетях потеря даже нескольких нейронов критична. Leaky ReLU выдает малое отрицательное значение (обычно 0.01 от входа), градиент не нулевой, нейрон остается активным.
• Кодирование частот: Вместо прямой подачи UV-координат на сеть, сначала пропустите их через синусы и косинусы разных частот. Это повышает качество без роста вычислений. Сети плохо учат высокочастотные паттерны (детали, резкие переходы) из низкоразмерных входов. Кодируя как [sin(2πu), cos(2πu), sin(2πv), cos(2πv)], даем компоненты частот. Сеть лучше захватывает гладкие и детальные паттерны, полезно для пространственных входов вроде UV. Для не-пространственных — менее актуально.

// Кодирование частот для лучшего представления нейронных текстур
float4 encodeUV(float2 uv) {
    float4 encoded;
    encoded.x = sin(uv.x * 2.0 * 3.14159);
    encoded.y = cos(uv.x * 2.0 * 3.14159);
    encoded.z = sin(uv.y * 2.0 * 3.14159);
    encoded.w = cos(uv.y * 2.0 * 3.14159);
    return encoded;
}

// Улучшенная сеть с кодированием частот
[Differentiable]
float3 evaluateNetworkWithEncoding(NeuralNetwork net, float2 uv) {
    float4 encoded = encodeUV(uv); // Теперь 4D-вход вместо 2D
    float3 output = float3(0.0, 0.0, 0.0);
    for (int i = 0; i < 3; i++) {
        output[i] = net.biases[i];
        for (int j = 0; j < 4; j++) {
            output[i] += net.weights[i * 4 + j] * encoded[j];
        }
    }
    return smoothActivation(output);
}

Как аппаратное обеспечение ускоряет кооперативные векторы

Современные GPU имеют Tensor Cores для эффективных матричных умножений. Но для Tensor Cores нужна кооперативная работа потоков.

Кооперативные векторы упрощают доступ к этому. Позволяют писать код как обычное умножение матрицы на вектор, компилятор сам мапит на Tensor Cores без упаковки или унифицированного потока.

Вот использование кооперативных векторов для ускорения сетей:

struct FeedForwardLayer<int InputSize, int OutputSize> {
    ByteAddressBuffer weights;
    uint weightsOffset;
    ByteAddressBuffer biases;
    uint biasesOffset;

    CoopVec<float, OutputSize> eval(CoopVec<float, InputSize> input) {
        let output = coopVecMatMulAdd<float, OutputSize>(
            input,
            CoopVecComponentType.Float32, // формат входа
            weights, weightsOffset,
            CoopVecComponentType.Float32, // формат весов
            biases, biasesOffset,
            CoopVecComponentType.Float32, // формат смещений
            CoopVecMatrixLayout.ColumnMajor, // layout матрицы
            false, // транспонирована ли матрица
            sizeof(float) * InputSize); // шаг матрицы

        return max(CoopVec<float, OutputSize>(0.0f), output); // ReLU
    }
}

Применения в реальности: что можно создать?

Описанные техники — база для множества применений нейронного шейдинга. Вот перспективные направления:

Нейронное сжатие текстур (NTC)

Нейронное сжатие текстур — одно из самых практичных применений. Традиционные форматы вроде BC1 и BC7 имеют пределы, но NTC дает выше качество при той же компрессии или лучше сжатие при равном качестве.

Идея — маленькая сеть как декодер, на вход низкоточностные латентные текстуры и позиционное кодирование. Преимущества:

• Переменный битрейт: Разное число латентных текстур с низкой разрядностью дает диапазон от 0.5 до 20 бит на пиксель.
• Независимый декодинг: Каждый пиксель декодируется отдельно, подходит для сэмплинга в шейдерах.
• Без галлюцинаций: Малые сети, обученные с нуля на текстуру, не генерируют артефакты вроде лишних деталей.

Детали и примеры в библиотеке NVIDIA для нейронного сжатия текстур.

Нейронные материалы

Нейронные материалы — мощное применение: учим сложные слоистые материалы и дистиллируем в компактные сети, быстрее оригинального шейдера.

Сеть принимает направление света, взгляда и латентные коды, выдает цвет материала. Для пространственной вариации — текстура латентных кодов как допвход.

Ключ — энкодер во время обучения переводит оригинальные текстуры в латентные, затем бейк для runtime. Масштабируется до 4K+ без проблем с оптимизацией по пикселям.

За пределами текстур и материалов

Принципы нейронного шейдинга применимы шире. Можно использовать для:

• Расчетов освещения: Аппроксимация сложных моделей быстрыми нейронными.
• Постобработки: Оптимальный тон-маппинг, цветокоррекция или стилизация.
• Обработки геометрии: Процедурная генерация или модификация геометрии.
• Анимации: Плавные интерполяции или процедурные движения.
• Процедурной генерации: Алгоритмическое создание контента с изученными паттернами.
• Деноизинга для рейтрейсинга: Снижение шума в рейтрейс-изображениях.
• Сжатия анимации: Эффективное сжатие данных анимации.
• Упрощения мешей: Упрощение 3D-мешей с сохранением вида.

Суть — находить дорогие вычисления или сложные связи, где нейронная аппроксимация поможет.

Будущее: почему нейронный шейдинг важен

Нейронный шейдинг — не просто техника, а смена парадигмы в графике реального времени. Обучаемые шейдеры открывают новые горизонты:

• Баланс качества и производительности: Сети настраиваются плавно под уровни качества, для естественных LOD.
• Расширяемость: Легко добавлять компоненты для сэмплинга или фильтрации.
• Гибкость платформ: Одни нейронные ассеты работают на разном железе — инференс на мощном, транскодинг на слабом.

Успех требует инструментария оптимизаций и отладки для нейронного шейдинга. Как с любой новинкой, есть кривая обучения, но отдача велика.

Первые шаги: что делать дальше

Экосистема нейронного шейдинга развивается быстро. Вот ключевые инструменты:

• Slang: Основной язык шейдинга с autodiff.
• SlangPy: Интерфейс Python для прототипов.
• SDK RTX Neural Shaders: Библиотека для инференса и обучения сетей.
• Примеры кооперативных векторов: Полные примеры ускорения.

Для быстрого старта с Slang и autodiff попробуйте в браузере Slang Playground. Для ресурсов — NVIDIA RTX Kit с поддержкой нейронного шейдинга. Глубже в концепции — курс NVIDIA по нейронному шейдингу на SIGGRAPH.

Технология готова для производства. Для сжатия текстур, материалов или новых идей нейронный шейдинг — мощный инструмент для качества и скорости в графике реального времени.