forked from denosaurs/netsaur
-
Notifications
You must be signed in to change notification settings - Fork 0
/
mod.ts
117 lines (104 loc) · 2.8 KB
/
mod.ts
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
import { dlopen, FetchOptions } from "../../../deps.ts";
import { GPUBackend } from "./backend.ts";
import { NoBackendError } from "../../core/api/error.ts";
import { BackendLoader, Engine } from "../../core/engine.ts";
import {
Backend,
BackendType,
Cost,
NetworkConfig,
SchedulerType,
} from "../../core/types.ts";
import { Sequential } from "../../core/mod.ts";
const options: FetchOptions = {
name: "netsaur_gpu",
url: new URL(import.meta.url).protocol !== "file:"
? new URL(
"https://github.com/denosaurs/netsaur/releases/download/0.2.14/",
import.meta.url,
)
: "./target/release/",
cache: "reloadAll",
};
const symbols = {
ffi_backend_create: {
parameters: ["buffer", "usize", "pointer"],
result: "usize",
} as const,
ffi_backend_train: {
parameters: ["usize", "buffer", "usize", "buffer", "usize"],
result: "void",
} as const,
ffi_backend_predict: {
parameters: ["usize", "buffer", "buffer", "usize", "buffer"],
result: "void",
} as const,
ffi_backend_save: {
parameters: ["usize", "pointer"],
result: "void",
} as const,
ffi_backend_load: {
parameters: ["buffer", "usize", "pointer"],
result: "usize",
} as const,
};
export type Library = Deno.DynamicLibrary<typeof symbols>;
export class GPUInstance {
static library?: Library;
static initialized = false;
static async init(silent = false) {
if (GPUInstance.initialized) return true;
GPUInstance.library = await dlopen(options, symbols);
GPUInstance.initialized = true;
if (!silent) console.log("GPU Backend Initialized");
return true;
}
}
export class GPUBackendLoader implements BackendLoader {
backend?: GPUBackend;
isSupported(): boolean {
return Deno.dlopen !== undefined;
}
async setup(silent = false) {
Engine.type = BackendType.GPU;
return await GPUInstance.init(silent);
}
loadBackend(config: NetworkConfig): Backend {
if (!GPUInstance.initialized) {
throw new NoBackendError(BackendType.GPU);
}
return this.backend
? this.backend
: GPUBackend.create(config, GPUInstance.library!);
}
load(buffer: Uint8Array): Sequential {
this.backend = GPUBackend.load(buffer, GPUInstance.library!);
const net = new Sequential({
size: [0],
layers: [],
cost: Cost.MSE,
scheduler: {
type: SchedulerType.None,
},
});
this.backend = undefined;
return net;
}
loadFile(path: string): Sequential {
this.backend = GPUBackend.loadFile(path, GPUInstance.library!);
const net = new Sequential({
size: [0],
layers: [],
cost: Cost.MSE,
scheduler: {
type: SchedulerType.None,
},
});
this.backend = undefined;
return net;
}
}
/**
* GPU Backend written in Rust.
*/
export const GPU = new GPUBackendLoader();