Решаем судоку на pytorch

Можно ли делать нейросети без обучения? Без кучи тестовых примеров?

Что обычно представляется, когда решается задача при помощи нейросетей? Формирование структуры сети. Написание loss функции. Куча примеров для обучения. Но в данном случае попробуем пойти другим путём. Воспользуемся нейросетью как языком программирования. Тем более что pytorch крайне прост и казуален. Задача кстати крайне неудобная для нейросетей. Обычно её решают перебором вариантов в той либо иной степени. А нейросети подобное делать тяжеловато. Но! Тем интереснее посмотреть на результаты.

Имеет ли это практический смысл? Да имеет. Можно получившуюся модель сконвертировать в onnx формат, а потом переконвертировать для запуска на нестандартном NPU. Например, на Maix III или Kendryte K210. Поэтому при написании нейросети будем поглядывать на список поддерживаемых Maix III операторов.

Что будет делать наша нейросеть? Решать судоку как человек, за множество проходов, убирая невозможные варианты. А если уж совсем тяжёлое судоку, то будет делать перебор различных вариантов постепенно отметая неверные.

Прежде всего глянем, а может уже до меня всё сделали? Смотрим, человек даже написал отчет stanford.edu. Результаты правда у него откровенно отвратительные. Решает только 86% ооочень-очень простых судоку, уровня показанных на первой картинке. Причём на каждое судоку он тратит 0.2 секунды на моём GeForce 3060. Получается классический вариант: А давайте возьмём нейросеть побольше, дадим ей миллион тестовых примеров и пускай нейросетка сама учится. Причём training accuracy=99.7% это скорее минус, а не плюс. Т.е. нейросетка как тупой ученик запомнила миллион тестовых примеров, но не поняла правил по которым всё это работает.

Поэтому придётся таки самим писать код.

Пишем код используя нейросети

Тут будет много технических деталей, которые не факт, что будут понятны без просмотра репозитория с кодом .

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

Смотрим на доступный нам инструментарий. На видеокарте от Nvidia поддерживаются все операторы, но мы же потом надеемся запустить нейросетку на чём-либо встраиваемом, поэтому смотрим на список операторов для Maix III. Conv2D и ReLu поддерживается. Depthwise min, max, abs, clamp, mul, add, sub так-же поддерживаются. Это удобные примитивы для создания fuzzy logic. MaxPool2D так-же поддерживается. А вот MaxUnpool2d к сожалению нет. Впрочем это и понятно, операция достаточно нестандартная, хоть и крайне полезная. Позже ещё вернусь к вопросу отсутствия MaxUnpool2d.

Из всего этого получается достаточно полный язык программирования. Если мы хотим делать логические операции в пределах одинаковых координат c,y,x то используем depthwise операции. nn.ReLu теоретически можно было бы и не пользоваться, т.к. это аналог torch.maximum(v, 0).

Если же мы хотим при вычислениях использовать данные с разным c,x,y, то нам уже придётся пользоваться операциями Conv2D/UpsamplingNearest2d и MaxPool2d/MaxUnpool2d.

А теперь попробуем попрограммировать. Сразу скажу - уже всё сделал и результаты положил в репозиторий на github. Поэтому дальше очень кратенько пробегусь по самым заметным моментам. Входные данные пускай представляются классически (как представляется 2D изображение в pytorch), в виде тензора размерностью M[N,C,H,W], где:

• N количество элементов (обычно это разные картинки одинакового размера).

• C (color/channels) количество цветов или каналов с информацией.

• H (height) высота изображения

• W (width) ширина изображения

В нашем случае это превращается в:

• N количество параллельно решаемых судоку

• С = 9 каналов, символизирующих возможность появления числа в этой ячейке.

• H,W = 9 так и остаются шириной и высотой поля судоку 9x9

В каждой из ячеек может храниться либо 0 либо 1. Изначально все пустые ячейки судоку заполнены единичками M[N,c,x,y]==1 для любого c, если в ячейке x,y у нас пусто. Если же к примеру в ячейке x,y точно известное число 3, то:

• M[N,c,x,y]==0 если c!=(3-1)

• M[N,c,x,y]==1 если c==(3-1)

Conv2D позволяет очень удачно делать логические операции со строками, столбцами и квадратами 3x3 (box). Хотим работать со строкой, используем kernel_size=(1, 9). Со столбцом - kernel_size=(9, 1). Для box надо использовать kernel_size=3, stride=3. Всё элементарно.

Это позволяет сделать нам базовые проверки:

• Если определённое число есть в какой либо ячейке этой строки, столбца или box, то значит в других ячейках его точно нет.

• Если число есть только в одной из ячеек, то значит оно может быть только в этой ячейке.

• Если число в строке или столбце находится только в одном из box, то значит в других box для этой строки/столбца оно не может находиться.

Этих условий уже достаточно, чтобы наш решатель судоку начал решать все судоку из puzzles0_kaggle, причём делать это всего за 4 прохода. И это происходит очень быстро. Впрочем, об этом расскажу, когда буду рассказывать о тестировании скорости. Т.е. уже на данном этапе у нас получился результат сильно лучше, чем получается при классическом "сделать непонятно какую структуру сети, дать ей кучу примеров и пусть сама разбирается". У нас получилась и 100% точность на решении простых судоку и скорость в миллион раз быстрее, чем классическая нейросетка.

Но мы же хотим большего верно? Решать medium, hard и extra hard судоку!

На medium уровне достаточно добавить правило, которое замечает двойки одинаковых чисел. Числа 1 и 8 подчёркнутые красным, могут находиться только в этих ячейках. Соответственно в остальных ячейках этого столбца они находиться не могут.

Правило получилось дорогим, очень дорогим. Я перебираю все комбинации возможных двоек чисел в ячейке. А это - 36 комбинаций. Можно глянуть на класс SudokuDigitsDoubles, как это происходит. Так-что это один из важных из кандидатов на оптимизацию. Но пока, к сожалению не придумал, как это сделать красиво.

Рекурсия

Дальше идут сложные и очень сложные судоку. Тут уже нужна рекурсия. Тут я реально долго думал. Пока не понял, как сделать стек внутри нейросети. Идея стека находится в классе Iterate2D в файле recursion_fun.py Этот класс выбирает последовательно числа из двухмерной матрицы и возвращает матрицу, в которой весь двухмерный тензор нулевой и есть только одна единичка - выбранный нами элемент. Причём если есть два элемента с одинаковым значением, то сначала будет выбран один элемент, а на следующем шаге итерации - другой. Вся магия заключается в MaxUnpool2d. Т.е сначала MaxPool2d применённый ко всему 2D изображению находит максимальный элемент. А потом MaxUnpool2d выбирает один и только один элемент и воссоздаёт маску обратно. Потом, пользуясь этой маской, можно будет занулить этот элемент, и соответственно в следующей итерации будет выбран следующий. Подобного поведения достаточно сразу для двух элементов рекурсии. Первый - нам нужно выбрать ячейку, в которой мы будем предполагать одно из нескольких значений для судоку. Второй - нам нужно организовать стек, для случая, если наше предположение неверное, чтобы восстановить исходное значение тензора.

Сама идея не очень сложная, но повозиться пришлось изрядно. Всё-таки у нас получается язык программирования, который местами даже более низкоуровневый, чем ассемблер. К тому-же с массовым параллелизмом (кто писал на Verilog, тот знает насколько при этом изменяется программирование).

Причём из-за отсутствия ветвлений нам приходится параллельно добавлять данные в стек и извлекать данные из него, или ничего не делать и только в конце выбирать, а что-же мы собственно хотим сделать? На этот момент можно глянуть в конце класса SudokuRecursionControl, там где несколько NNSelect расположено.

А вот вам и последовательность решения сложного судоку. Красными маленькими квадратиками отмечены числа, которые отложены "на потом". Если текущая ветка рекурсии окажется неправильной.

Производительность

И хоть этот код и не претендует на рекорды производительности, но всё-же интересно сравнить с лучшими. В качестве лучшего у нас будет выступать tdoku. Эта штука умудряется решить сложные судоку буквально за несколько микросекунд. А puzzles0_kaggle оно называет "почти решёнными судоку". Напомню - нейросеть собранная по классическим рецептам даже с puzzles0_kaggle не особо хорошо ладит.

Нейросеть на pytorch бессмысленно заставлять решать судоку по одному. Потому как каждый запуск нейросети - это минимум пара миллисекунд времени. Поэтому приходится давать на параллельное решение минимум по 10 тыс. судоку. Это несколько непрактично, но таки позволяет замерить скорость работы.

Запускаем на GeForce 3060. Результаты получились такие:

Если сравнивать с tdoku, который оптимизирован под SIMD операции и использует значительно более продвинутое бинарное представление судоку, то результаты сильно хуже. Но если учесть, что у нас расчёты во float, то "всё не так плохо". Если бы у нас была возможность выкидывать уже решённые судоку из тензора и заменять их новыми, то вполне возможно была-бы сравнимая производительность. Т.е. тут мы упираемся в то, что хотя-бы немножко условных переходов желательно иметь при решении этой задачи. А ещё можно попытаться перейти на int8 точность расчётов, что тоже может дать прирост в производительности в несколько раз.

Отсутствие MaxUnpool2d - это так-же не фатальная проблема. В случае, если все ячейки имеют разные значения, она работает так-же как UpsamplingNearest2d с последующим сравнением с максимальным элементом. А уникальность ячеек мы можем обеспечить добавляя небольшое число к каждой из ячеек так, что-бы все они были разными. Сложность тут только в том, что-бы не вылезти за диапазон int8. Здесь мы так-же упираемся в заточенность софта для квантизации под процесс обучения. Типичный софт выглядит так - мы ему даём нейросеть в onnx формате и кучу входных примеров. А он уже сам определяет, в каком диапазоне у нас числа находятся. Для ручного задания весов Conv2D, когда мы точно знаем диапазоны входных-выходных значений, это может только привести к неожиданным багам.

Вывод

Подобный подход имеет смысл. Таким способом можно получить как более быструю, так и более надёжную нейросеть. Получилось даже рекурсию реализовать, в чем был изначально совершенно не уверен. Так-же у нас есть возможность запускать код на NPU. На таком низком уровне конечно тяжеловато писать код, но возможно. Не удивлюсь, если в будущем напишут и более высокоуровневый язык программирования, который будет компилировать программу в нейросеть. Естественно этот язык не будет похож ни на один из языков, запускаемых на CPU.

PS: Ещё раз продублирую ссылку. Если хочется посмотреть на код, то репозиторий этого проекта находится на github.