初始化项目结构

.gitignore

@@ -0,0 +1,6 @@
+dataset/*
+checkpoints/*
+!dataset/.keep
+!checkpoints/.keep
+keys
+.idea

README.md

@@ -0,0 +1,31 @@
+# classification-models-tensorflow
+
+## 项目说明
+此项目包含AlexNet模型的Keras框架实现和VGG16模型的tensorflow框架实现和与其对应的模型训练文件
+
+## 项目文件说明
+```text
+classification-models-tensorflow
+    ├── README.md
+    ├── checkpoints  # 保存所有的权重信息
+    ├── export_onnx.py  # 模型权重转换为onnx脚本
+    ├── models  # 模型定义
+    │   └── AlexNet.py
+    ├── train_alexnet.py  # AlexNet模型训练脚本
+    └── train_vgg16.py  # VGG16模型训练脚本
+```
+
+## 训练命令
+- AlexNet
+```shell
+python train_alexnet.py --data-path dataset/imagenette2-320 --output-dir checkpoints/alexnet --batch-size 64 --epochs 90
+```
+- VGG16
+```shell
+python train_vgg16.py --data-path dataset/imagenette2-320 --output-dir checkpoints/vgg16 --batch-size 64 --epochs 90
+```
+
+## 模型训练权重转换为onnx
+```shell
+python export_onnx.py 
+```

checkpoints/.keep

@@ -0,0 +1 @@
+此目录用于存放模型训练权重信息

dataset/.keep

@@ -0,0 +1 @@
+此目录用于存放数据集

export_onnx.py

@@ -0,0 +1,17 @@
+import tensorflow as tf
+import tf2onnx
+
+
+# 1. 加载保存的 TensorFlow 模型
+model = tf.keras.models.load_model('saved_model')
+
+# 2. 定义模型的输入签名
+spec = (tf.TensorSpec((None, 32, 32, 3), tf.float32, name="input"),)
+
+# 3. 转换并保存为 ONNX 格式
+output_path = 'vgg16_tensorflow.onnx'
+model_proto, _ = tf2onnx.convert.from_keras(model, input_signature=spec, opset=11)
+
+# 4. 保存 ONNX 模型
+with open(output_path, "wb") as f:
+    f.write(model_proto.SerializeToString())

models/AlexNet.py

@@ -0,0 +1,40 @@
+from keras.models import Sequential
+from keras.layers import Conv2D, MaxPooling2D, Dense, Flatten, Dropout, ReLU, AveragePooling2D
+
+
+# 使用 Keras 重新定义 AlexNet 模型
+def create_model(input_shape=(224, 224, 3), num_classes=10):
+    model = Sequential([
+        # 特征提取部分
+        Conv2D(64, kernel_size=(11, 11), strides=(4, 4), padding="same", input_shape=input_shape),
+        ReLU(),
+        MaxPooling2D(pool_size=(3, 3), strides=(2, 2)),
+
+        Conv2D(192, kernel_size=(5, 5), padding="same"),
+        ReLU(),
+        MaxPooling2D(pool_size=(3, 3), strides=(2, 2)),
+
+        Conv2D(384, kernel_size=(3, 3), padding="same"),
+        ReLU(),
+
+        Conv2D(256, kernel_size=(3, 3), padding="same"),
+        ReLU(),
+
+        Conv2D(256, kernel_size=(3, 3), padding="same"),
+        ReLU(),
+        MaxPooling2D(pool_size=(3, 3), strides=(2, 2)),
+
+        # 全局平均池化部分
+        AveragePooling2D(pool_size=(6, 6)),
+
+        # 分类部分
+        Flatten(),
+        Dropout(0.5),
+        Dense(4096),
+        ReLU(),
+        Dropout(0.5),
+        Dense(4096),
+        ReLU(),
+        Dense(num_classes, activation='softmax')
+    ])
+    return model

train_alexnet.py

@@ -0,0 +1,139 @@
+import os
+
+import tensorflow as tf
+from keras.optimizers import Adam, SGD
+from keras.callbacks import ModelCheckpoint, CSVLogger
+from models.AlexNet import create_model
+from tensorflow.keras.preprocessing import image_dataset_from_directory
+
+
+
+def load_data(train_dir, val_dir, img_size=(224, 224), batch_size=32):
+    # Define data augmentation for the training set
+    train_datagen = tf.keras.Sequential([
+        tf.keras.layers.RandomFlip('horizontal'),
+        tf.keras.layers.RandomRotation(0.2),
+        tf.keras.layers.RandomZoom(0.2),
+        tf.keras.layers.RandomContrast(0.2),
+    ])
+
+    # Load training dataset
+    train_dataset = image_dataset_from_directory(
+        train_dir,
+        image_size=img_size,  # Resize images to (224, 224)
+        batch_size=batch_size,
+        label_mode='categorical',  # Return integer labels
+        shuffle=True
+    )
+
+    # Load validation dataset
+    val_dataset = image_dataset_from_directory(
+        val_dir,
+        image_size=img_size,  # Resize images to (224, 224)
+        batch_size=batch_size,
+        label_mode='categorical',  # Return integer labels
+        shuffle=False
+    )
+
+    # Normalize the datasets (rescale pixel values to [0, 1])
+    train_dataset = train_dataset.map(
+        lambda x, y: (train_datagen(x) / 255.0, y),
+    )
+
+    val_dataset = val_dataset.map(
+        lambda x, y: (x / 255.0, y),
+    )
+
+    return train_dataset, val_dataset
+
+
+def train_model(args, train_data, val_data):
+    # Create model
+    model = create_model()
+
+    # 调整学习率
+    learning_rate = args.lr if args.lr else 1e-2
+    # optimizer = SGD(learning_rate=learning_rate, momentum=args.momentum)
+
+    # Compile model
+    model.compile(optimizer=Adam(learning_rate=0.0001), loss='categorical_crossentropy', metrics=['accuracy'])
+
+    # Check if a checkpoint exists and determine the initial_epoch
+    latest_checkpoint = tf.train.latest_checkpoint(args.output_dir)
+    if latest_checkpoint:
+        initial_epoch = int(latest_checkpoint.split('_')[-1].split('.')[0])  # Get the last epoch from filename
+        print(f"Resuming training from epoch {initial_epoch}")
+    else:
+        initial_epoch = 0
+
+    # Define CSVLogger to log training history to a CSV file
+    csv_logger = CSVLogger(os.path.join(args.output_dir, 'training_log.csv'), append=True)
+
+    # Define ModelCheckpoint callback to save weights for each epoch
+    checkpoint_callback = ModelCheckpoint(
+        os.path.join(args.output_dir, 'alexnet_{epoch:03d}.h5'),  # Save weights as alexnet_{epoch}.h5
+        save_weights_only=True,
+        save_freq='epoch',  # Save after every epoch
+        verbose=1
+    )
+
+    # Train the model
+    history = model.fit(
+        train_data,
+        epochs=args.epochs,
+        validation_data=val_data,
+        initial_epoch=initial_epoch,
+        callbacks=[csv_logger, checkpoint_callback],  # Add checkpoint callback
+    )
+
+    return history
+
+
+def get_args_parser(add_help=True):
+    import argparse
+
+    parser = argparse.ArgumentParser(description="PyTorch Classification Training", add_help=add_help)
+
+    parser.add_argument("--data-path", default="dataset/imagenette2-320", type=str, help="dataset path")
+    parser.add_argument("--output-dir", default="checkpoints/alexnet", type=str, help="path to save outputs")
+
+    parser.add_argument("--device", default="cuda", type=str, help="device (Use cuda or cpu Default: cuda)")
+    parser.add_argument(
+        "-b", "--batch-size", default=64, type=int, help="images per gpu, the total batch size is $NGPU x batch_size"
+    )
+    parser.add_argument("--epochs", default=90, type=int, metavar="N", help="number of total epochs to run")
+
+    parser.add_argument("--opt", default="sgd", type=str, help="optimizer")
+    parser.add_argument("--lr", default=0.1, type=float, help="initial learning rate")
+    parser.add_argument("--momentum", default=0.9, type=float, metavar="M", help="momentum")
+    parser.add_argument("--lr-scheduler", default="steplr", type=str, help="the lr scheduler (default: steplr)")
+    parser.add_argument("--lr-warmup-epochs", default=0, type=int, help="the number of epochs to warmup (default: 0)")
+    parser.add_argument(
+        "--lr-warmup-method", default="constant", type=str, help="the warmup method (default: constant)"
+    )
+    parser.add_argument("--lr-warmup-decay", default=0.01, type=float, help="the decay for lr")
+    parser.add_argument("--lr-step-size", default=30, type=int, help="decrease lr every step-size epochs")
+    parser.add_argument("--lr-gamma", default=0.1, type=float, help="decrease lr by a factor of lr-gamma")
+    parser.add_argument("--lr-min", default=0.0, type=float, help="minimum lr of lr schedule (default: 0.0)")
+    parser.add_argument("--start-epoch", default=0, type=int, metavar="N", help="start epoch")
+
+    parser.add_argument(
+        "--input-size", default=224, type=int, help="the random crop size used for training (default: 224)"
+    )
+    return parser
+
+if __name__ == "__main__":
+    args = get_args_parser().parse_args()
+
+    # Set directories for your custom dataset
+    train_dir = os.path.join(args.data_path, "train")
+    val_dir = os.path.join(args.data_path, "val")
+
+    # Set the directory where you want to save weights
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    # Load data
+    train_data, val_data = load_data(train_dir, val_dir, img_size=(args.input_size, args.input_size), batch_size=args.batch_size)
+
+    # Start training
+    train_model(args, train_data, val_data)