# gans-jax

**Repository Path**: fmscole/gans-jax

## Basic Information

- **Project Name**: gans-jax
- **Description**: No description available
- **Primary Language**: Unknown
- **License**: Unlicense
- **Default Branch**: master
- **Homepage**: None
- **GVP Project**: No

## Statistics

- **Stars**: 0
- **Forks**: 0
- **Created**: 2026-05-14
- **Last Updated**: 2026-06-03

## Categories & Tags

**Categories**: Uncategorized

**Tags**: None

## README

# gans-jax

基于 JAX/Stax 和 PyTorch 的深度学习模型实现，涵盖生成模型、Transformer、目标检测等多个方向。

## 项目结构

```
gans-jax/
├── src/
│   ├── lib/              # 共享工具库（stax_plus 扩展层、数据集加载）
│   ├── gans/             # GAN 变体（vanilla、WGAN、CGAN、InfoGAN）
│   ├── diffusion/        # 扩散模型（DDPM + UNet）
│   ├── transformers/     # Transformer 实现
│   ├── translation/      # Seq2seq 翻译模型（GRU/LSTM + 注意力）
│   ├── flash_attention/  # Flash Attention（JAX 自定义 VJP）
│   ├── vit/              # Vision Transformer 变体
│   ├── yolov5/           # YOLOv5 目标检测
│   └── vae/              # 变分自编码器 / 稀疏自编码器
├── notebooks/            # Jupyter Notebook（GAN、Diffusion）
├── scripts/              # 训练与工具脚本
├── docs/                 # 论文和数据集文档
├── data/                 # 数据集（MNIST、CIFAR-10、平行语料）
├── LICENSE
├── .gitignore
└── README.md
```

## 各模块说明

### GAN（`src/gans/`）
基础 GAN、WGAN、条件 GAN、InfoGAN，分别用 Flax、Stax（JAX）和 PyTorch 实现。

### Diffusion（`src/diffusion/`）
DDPM + UNet 骨干网络，支持 MNIST 图像生成。包含 Flax、Stax、PyTorch 三种实现。

### Transformer（`src/transformers/`）
基于 Stax 的 Transformer 编码器/解码器，集成 Flash Attention，包含时间序列变体。

### 翻译模型（`src/translation/`）
基于 GRU/LSTM 编码器-解码器的神经机器翻译，支持注意力机制、Flash Attention 和 KV-Cache 推理。

### Flash Attention（`src/flash_attention/`）
IO 感知的精确注意力计算，JAX 自定义前向/反向传播（custom VJP），支持因果掩码和填充掩码。

### Vision Transformer（`src/vit/`）
面向 MNIST、CIFAR-10、ImageNet 的 ViT 实现，Stax 和 PyTorch 双版本，含 Flash Attention 变体。

### YOLOv5（`src/yolov5/`）
JAX/Stax 实现的 YOLOv5 目标检测器，支持 VOC 和 COCO 数据集。

### VAE（`src/vae/`）
变分自编码器（Stax）和稀疏自编码器（PyTorch），基于 MNIST 训练。

## 依赖

- JAX / Flax / Optax
- PyTorch / torchvision
- NumPy / Matplotlib

## License

MIT