Adept | Машинное обучение | Открытое ПО
@channel_adept_platform
Сделал детальное представление вычислительного графа и теперь его можно преобразовать в другое IR.
Вот например для такого небольшого перцептрона:
class MLP : public Module {
public:
MLP()
: linear1(28 * 28, 1024),
linear2(1024, 512),
linear3(512, 256),
linear4(256, 10) {
...
}
Tensor forward(Tensor input) {
auto out = l1(input);
out = relu(out);
out = l2(out);
out = relu(out);
out = l3(out);
out = relu(out);
out = l4(out);
return out;
}
...
};
Получается такое представление в StableHLO формате:
module {
func.func @main(%arg0: tensor<784x784xf32>, %arg1: tensor<1024x784xf32>, %arg2: tensor<1024xf32>, %arg3: tensor<512x1024xf32>, %arg4: tensor<512xf32>, %arg5: tensor<256x512xf32>, %arg6: tensor<256xf32>, %arg7: tensor<10x256xf32>, %arg8: tensor<10xf32>) -> tensor<784x10xf32> {
%0 = stablehlo.transpose %arg1, permutation = [1, 0] : (tensor<1024x784xf32>) -> tensor<784x1024xf32>
%1 = stablehlo.dot_general %arg0, %0, contracting_dims = [1] x [0] : (tensor<784x784xf32>, tensor<784x1024xf32>) -> tensor<784x1024xf32>
%2 = stablehlo.broadcast_in_dim %arg2, dims = [1] : (tensor<1024xf32>) -> tensor<784x1024xf32>
%3 = stablehlo.add %1, %2 : tensor<784x1024xf32>
%4 = stablehlo.constant dense<0> : tensor<784x1024xf32>
%5 = stablehlo.maximum %3, %4 : tensor<784x1024xf32>
%6 = stablehlo.transpose %arg3, permutation = [1, 0] : (tensor<512x1024xf32>) -> tensor<1024x512xf32>
%7 = stablehlo.dot_general %5, %6, contracting_dims = [1] x [0] : (tensor<784x1024xf32>, tensor<1024x512xf32>) -> tensor<784x512xf32>
%8 = stablehlo.broadcast_in_dim %arg4, dims = [1] : (tensor<512xf32>) -> tensor<784x512xf32>
%9 = stablehlo.add %7, %8 : tensor<784x512xf32>
%10 = stablehlo.constant dense<0> : tensor<784x512xf32>
%11 = stablehlo.maximum %9, %10 : tensor<784x512xf32>
%12 = stablehlo.transpose %arg5, permutation = [1, 0] : (tensor<256x512xf32>) -> tensor<512x256xf32>
%13 = stablehlo.dot_general %11, %12, contracting_dims = [1] x [0] : (tensor<784x512xf32>, tensor<512x256xf32>) -> tensor<784x256xf32>
%14 = stablehlo.broadcast_in_dim %arg6, dims = [1] : (tensor<256xf32>) -> tensor<784x256xf32>
%15 = stablehlo.add %13, %14 : tensor<784x256xf32>
%16 = stablehlo.constant dense<0> : tensor<784x256xf32>
%17 = stablehlo.maximum %15, %16 : tensor<784x256xf32>
%18 = stablehlo.transpose %arg7, permutation = [1, 0] : (tensor<10x256xf32>) -> tensor<256x10xf32>
%19 = stablehlo.dot_general %17, %18, contracting_dims = [1] x [0] : (tensor<784x256xf32>, tensor<256x10xf32>) -> tensor<784x10xf32>
%20 = stablehlo.broadcast_in_dim %arg8, dims = [1] : (tensor<10xf32>) -> tensor<784x10xf32>
%21 = stablehlo.add %19, %20 : tensor<784x10xf32>
return %21 : tensor<784x10xf32>
}
}
Далее его надо преобразовать тензорным компиляторм. Эксперименты продолжаются.
Если у вас установлено приложение,
вы можете сразу перейти в канал
вы можете сразу перейти в канал