{"id":29530100,"url":"https://github.com/duyongan/sunstreaker","last_synced_at":"2025-09-12T21:24:42.298Z","repository":{"id":63393041,"uuid":"567243313","full_name":"duyongan/sunstreaker","owner":"duyongan","description":"以jax为后端的类似keras的框架","archived":false,"fork":false,"pushed_at":"2023-01-13T06:37:29.000Z","size":511,"stargazers_count":98,"open_issues_count":0,"forks_count":3,"subscribers_count":2,"default_branch":"main","last_synced_at":"2025-06-07T06:03:31.915Z","etag":null,"topics":["beginner-friendly","data-science","deep-learning","deep-learning-algorithms","deep-learning-framework","deep-learning-library","deep-learning-tutorial","deep-neural-networks","jax","keras","machine-learning","ml","neural-network","nlp","numpy","python","pytorch","scikit-learn","tensorflow"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"apache-2.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/duyongan.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null}},"created_at":"2022-11-17T11:33:38.000Z","updated_at":"2025-05-01T12:30:40.000Z","dependencies_parsed_at":"2023-02-09T15:01:37.576Z","dependency_job_id":null,"html_url":"https://github.com/duyongan/sunstreaker","commit_stats":null,"previous_names":[],"tags_count":1,"template":false,"template_full_name":null,"purl":"pkg:github/duyongan/sunstreaker","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/duyongan%2Fsunstreaker","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/duyongan%2Fsunstreaker/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/duyongan%2Fsunstreaker/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/duyongan%2Fsunstreaker/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/duyongan","download_url":"https://codeload.github.com/duyongan/sunstreaker/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/duyongan%2Fsunstreaker/sbom","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":265419639,"owners_count":23761848,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["beginner-friendly","data-science","deep-learning","deep-learning-algorithms","deep-learning-framework","deep-learning-library","deep-learning-tutorial","deep-neural-networks","jax","keras","machine-learning","ml","neural-network","nlp","numpy","python","pytorch","scikit-learn","tensorflow"],"created_at":"2025-07-16T22:02:11.100Z","updated_at":"2025-07-16T22:06:53.371Z","avatar_url":"https://github.com/duyongan.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"\u003cdiv align=\"center\"\u003e\n    \u003cimg src=\"img/logo.png\" alt=\"logo\" style=\"zoom:30%\"/\u003e\n\u003c/div\u003e\n\n# Sunstreaker  \n\n源码清晰明了，使用简单好搞\n\n\n## 目标\n\n- [x] 源码清晰简洁，利于算法学习与实验\n- [x] 快速实验新改进想法\n- [x] 快速复现新论文\n- [ ] 快速分布式训练一个大模型\n- [ ] 快速使用开源模型权重\n\n\n## 说明\n\n* 本项目采用小步快走的形式，欢迎start，但不建议fork，因为\u003cfont color=\"red\" size=4\u003e更新速度比较快\u003c/font\u003e。\n* 本项目用于学习与实验，切勿用于生产\n\n## 欢迎关注公众号：无数据不智能\n\n\u003cdiv align=\"center\"\u003e\n    \u003cimg src=\"img/blog.jpg\" alt=\"logo\" style=\"zoom:50%\"/\u003e\n\u003c/div\u003e\n\n## 安装\n\u003e tensorflow只是加载demo数据需要，也可以不装\n\n\n\n### windows \n\n1. 安装jax\n   * cpu \n   ```\n   pip install jax[cpu]==0.3.14 -f https://whls.blob.core.windows.net/unstable/index.html --use-deprecated legacy-resolver\n   ```\n   * gpu\n   ```\n   pip install jax[cuda111]==0.3.14 -f https://whls.blob.core.windows.net/unstable/index.html --use-deprecated legacy-resolver\n   ```\n2. 安装Graphviz\n   * [exe安装下载](http://graphviz.org/download/)\n   * pygraphviz\n   ```\n    pip install --global-option=build_ext `\n                  --global-option=\"-IC:\\Program Files\\Graphviz\\include\" `\n                  --global-option=\"-LC:\\Program Files\\Graphviz\\lib\" `\n                  pygraphviz\n   ```\n3. pip install -r requirements.txt\n4. pip install sunstreaker \n### linux\n1. 安装jax\n\n   - cpu\n\n   ```\n   pip install --upgrade jax[cpu]==0.3.14\n   ```\n\n   - gpu\n\n   ```\n   pip install --upgrade jax[cuda]==0.3.14 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\n   ```\n\n2. [安装Graphviz](http://graphviz.org/download/)\n\n3. pip install -r requirements.txt\n\n4. pip install sunstreaker \n\n\n\n\n\n## 使用\n\n\n\n### 用tensorflow_datasets搞些数据\n\n```python\nimport tensorflow as tf\n\ntf.config.set_visible_devices([], 'GPU')\nimport math\nimport asyncio\nimport tensorflow_datasets as tfds\nfrom sunstreaker.data import Dataloader\nfrom sunstreaker.layers import Flatten\nfrom sunstreaker.layers.activations import Softmax\nimport jax.numpy as jnp\nfrom sunstreaker.losses import categorical_crossentropy\nfrom sunstreaker.metrics import categorical_accuracy\nfrom sunstreaker.optimizers import RMSProp\n\n\ndef load(batch_size: int, func):\n   async def tfds_load_data() -\u003e Dataloader:\n      ds, info = tfds.load(name=\"mnist\", split=[\"train\", \"test\"], as_supervised=True, with_info=True,\n                           shuffle_files=True, batch_size=batch_size)\n      train_ds, valid_ds = ds\n      train_ds, valid_ds = func(train_ds), func(train_ds)\n      train_ds, valid_ds = train_ds.cache().repeat(), valid_ds.cache().repeat()\n      input_shape = tuple(list(info.features[\"image\"].shape))\n      num_train_batches = math.ceil(info.splits[\"train\"].num_examples / batch_size)\n      num_val_batches = math.ceil(info.splits[\"test\"].num_examples / batch_size)\n      return Dataloader(\n         train_data=iter(tfds.as_numpy(train_ds)), val_data=iter(tfds.as_numpy(valid_ds)),\n         input_shape=input_shape, batch_size=batch_size,\n         num_train_batches=num_train_batches, num_val_batches=num_val_batches\n      )\n\n   return asyncio.run(tfds_load_data())\n\n\ndef load_dataset(batch_size: int):\n   def func(ds):\n      return ds.map(lambda x, y: (tf.divide(tf.cast(x, dtype=tf.float32), 255.0), tf.one_hot(y, depth=10)))\n\n   return asyncio.run(load(batch_size, func))\n\n\ndef load_dataset_muti(batch_size: int):\n   def func(ds):\n      return ds.map(lambda x, y: ({\"img\": tf.divide(tf.cast(x, dtype=tf.float32), 255.0)}, {\"out1\": tf.one_hot(y, depth=10)}))\n\n   return asyncio.run(load(batch_size, func))\n```\n\n### 序贯式编程\n\n```python3\nfrom sunstreaker.engine.sequential import Model\n\ndata = load_dataset(batch_size=1024)\nmodel = Model([Input(input_shape=(28, 28, 1)), Flatten(), Dense(100), Dense(10), Softmax()])\n```\n\n### 函数式编程\n\n```python\ndata = load_dataset_muti(batch_size=1024)\ninputs = Input(input_shape=(28, 28, 1), name=\"img\")\nflatten = Flatten()(inputs)\ndense1 = Dense(100, activation='relu')(flatten)\ndense2 = Dense(10, use_bias=False)(dense1)\noutputs = Softmax(name=\"out1\")(dense2)\n\nfrom sunstreaker.engine.functional import Model\n\nmodel = Model(inputs=inputs, outputs=outputs)\n```\n\n\n\n\n### 当你是一个老手\n\n```python3\nfrom sunstreaker import Model\n\ndata = load_dataset(batch_size=1024)\n\nclass MyModel(Model):\n    def build(self, rng=None):\n        self.W = self.add_weight((784, 10))\n        self.flatten = Flatten()\n        self.softmax = Softmax()\n        return (10,), [(self.W,)]\n\n    def call(self, params, inputs, trainable=True, **kwargs):\n        self.W, = params[0]\n        x = self.flatten.forward(params=[], inputs=inputs)\n        x = jnp.dot(x, self.W)\n        y = self.softmax.forward(params=[], inputs=x)\n        return y\n\n\nmodel = MyModel()\n```\n\n### 编译、训练、保存\n\n```python3\nmodel.compile(loss=categorical_crossentropy, optimizer=RMSProp(lr=0.001), metrics=[categorical_accuracy])\nmodel.fit(data, epochs=10)\nmodel.save(\"tfds_mnist_v2\")\n```\n\n### 模型结构打印\n\n```python3\nmodel.summary()\nmodel.plot_model() \n```\n\n```commandline\n+--------+-----------+---------+-------------+--------------+\n| number | name      | class   | input_shape | output_shape |\n+--------+-----------+---------+-------------+--------------+\n| 0      | input_0   | Input   | (28, 28, 1) | (28, 28, 1)  |\n| 1      | flatten_1 | Flatten | (28, 28, 1) | (784,)       |\n| 2      | dense_2   | Dense   | (784,)      | (100,)       |\n| 3      | dense_4   | Dense   | (100,)      | (10,)        |\n| 4      | softmax_6 | Softmax | (10,)       | (10,)        |\n+--------+-----------+---------+-------------+--------------+\n```\n\n\u003cdiv align=\"left\"\u003e\n    \u003cimg src=\"img/model.png\" alt=\"logo\" /\u003e\n\u003c/div\u003e\n\n### 损失与评价可视化\n\n```python\nmodel.plot_losses()\nmodel.plot_accuracy()\n```\n\u003cdiv align=\"left\"\u003e\n    \u003cimg src=\"img/losses.png\" alt=\"logo\" /\u003e\n\u003c/div\u003e\n\u003cdiv align=\"left\"\u003e\n    \u003cimg src=\"img/accuracy_categorical_accuracy.png\" alt=\"logo\" /\u003e\n\u003c/div\u003e\n\n\n## 功能\n\n### 0.0.1.dev更新\n\n| activations |  layers   |             losses             |              metrics              | optimizers |\n| :---------: | :-------: | :----------------------------: | :-------------------------------: | :--------: |\n|   Linear    |   Dense   |      binary_crossentropy       |          binary_accuracy          |    SGD     |\n|   Softmax   |  Flatten  |    categorical_crossentropy    |             accuracy              |    SM3     |\n|    Relu     |  Dropout  |       mean_squared_error       |       categorical_accuracy        |  Adagrad   |\n|   Sigmoid   |  Conv2D   |      mean_absolute_error       |    sparse_categorical_accuracy    |    Adam    |\n|     Elu     | MaxPool2D | mean_squared_logarithmic_error |    cosine_similarity_accuracy     |   Adamax   |\n|  LeakyRelu  | AveragePooling2D |             hinge              |    top_k_categorical_accuracy     |  RMSProp   |\n|    Gelu     |    GRU    |         kl_divergence          | sparse_top_k_categorical_accuracy |    FTRL    |\n|             |           |             huber              |                                   |            |\n\n### 0.0.2.dev更新\n\n|           layers           |  losses  |\n| :------------------------: | :------: |\n|         Embedding          | l2_error |\n|           Lambda           |          |\n|            Add             |          |\n|        Concatenate         |          |\n|            Dot             |          |\n|          Multiply          |          |\n|     LayerNormalization     |          |\n|   InstanceNormalization    |          |\n|     BatchNormalization     |          |\n|     GroupNormalization     |          |\n| LocalResponseNormalization |          |\n|        UpSampling2D        |          |\n\n### 0.0.3.dev更新\n\n|  initializer   | activations |\n| :------------: | :---------: |\n|     zeros      |    Swish    |\n|      ones      |             |\n|    constant    |             |\n|    uniform     |             |\n|     normal     |             |\n|   orthogonal   |             |\n|  LecunUniform  |             |\n|  LecunNormal   |             |\n|  GlorotNormal  |             |\n| GlorotUniform  |             |\n|    HeNormal    |             |\n|   HeUniform    |             |\n| KaimingUniform |             |\n| KaimingNormal  |             |\n|  XavierNormal  |             |\n| XavierUniform  |             |\n|    Identity    |             |\n\n\n### 0.0.4.dev更新\n\n**内核改动**\n\n1. Layer call 函数不再需要传入params，build输出不再需要输出params，以dense为例\n\n   ```python\n   class Dense(Layer):\n       def __init__(self, units, activation=None, use_bias=True, kernel_initializer=GlorotUniform(), bias_initializer=Zeros(), **kwargs):\n           super().__init__(**kwargs)\n           self.use_bias = use_bias\n           self.activation = activations.get(activation)()\n           self.units = int(units) if not isinstance(units, int) else units\n           self.kernel_initializer = kernel_initializer\n           self.bias_initializer = bias_initializer\n   \n       def build(self):\n           output_shape = self.input_shape[:-1] + (self.units,)\n           self.add_weight(\"kernel\", (self.input_shape[-1], self.units), initializer=self.kernel_initializer, seed=k1)\n           if self.use_bias:\n               self.add_weight(\"bias\", (self.units,), initializer=self.bias_initializer, seed=k2)\n           return output_shape\n   \n       def call(self, inputs, **kwargs):\n           kernel = self.get_weight(\"kernel\")\n           if self.use_bias:\n               bias = self.get_weight(\"bias\")\n               outputs = jnp.dot(inputs, kernel) + bias\n           else:\n               outputs = jnp.dot(inputs, kernel)\n           outputs = self.activation.forward(params=None, inputs=outputs)\n           return outputs\n   ```\n\n2. Model params变为有序字典，方便大模型参数加载\n\n3. build不再需要输入随机种子，由内核自动分配\n\n### 0.0.5.dev更新\n\n|    application    |       layers       |\n| :---------------: | :----------------: |\n| transformers/bert | MultiHeadAttention |\n|                   | PositionEmbedding  |\n|                   |    FeedForward     |\n|                   |    ScaleOffset     |\n|                   |     Activation     |\n\n\n### 0.0.6.dev更新\n\n|  application   | optimizers |\n| :------------: | :--------: |\n| diffusion/DDPM |   AdamW    |\n\n\n\n\n## 引用\n\n* https://github.com/google/jax\n* https://github.com/google/flax\n* https://github.com/keras-team/keras\n* https://github.com/umangjpatel/kerax\n* https://github.com/bojone/bert4keras\n* https://github.com/ddbourgin/numpy-ml\n* https://github.com/huggingface/transformers\n\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fduyongan%2Fsunstreaker","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fduyongan%2Fsunstreaker","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fduyongan%2Fsunstreaker/lists"}