87owo/PYAS
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add files via upload

Browse Source
This commit is contained in:
87owo authored and GitHub committed 2025-12-28 14:20:35 +08:00
1 parent 10ecb35cca
commit ef68a89b45
1 file changed
+31 -26
Download Patch File Download Diff File Expand all files Collapse all files
Engine/Model/train.py
+31 -26
View File
@@ -11,7 +11,7 @@ tf.keras.mixed_precision.set_global_policy('mixed_float16')
####################################################################################################
def load_dataset(data_dir, image_size=(224, 224), val_split=0.01, batch_size=256, color_mode="grayscale"):
def load_dataset(data_dir, image_size=(224, 224), val_split=0.00001, batch_size=256, color_mode="grayscale"):
channels = 1 if color_mode == "grayscale" else 3
file_paths, labels, class_indices = get_file_list(data_dir)
@@ -42,10 +42,7 @@ def load_dataset(data_dir, image_size=(224, 224), val_split=0.01, batch_size=256
def parse_fn(filename, label, image_size, channels, num_classes):
img = tf.io.read_file(filename)
img = tf.image.decode_image(img, channels=channels, expand_animations=False)
img = tf.cond(
tf.reduce_all(tf.equal(tf.shape(img)[:2], image_size)),
lambda: tf.cast(img, tf.float32),
lambda: tf.image.resize(img, image_size))
img = tf.cond(tf.reduce_all(tf.equal(tf.shape(img)[:2], image_size)), lambda: tf.cast(img, tf.float32), lambda: tf.image.resize(img, image_size))
img = tf.cast(img, tf.float32) / 255.0
label = tf.one_hot(label, depth=num_classes)
return img, label
@@ -55,6 +52,7 @@ def get_file_list(directory):
for i, cls in enumerate(sorted([d for d in os.listdir(directory) if os.path.isdir(os.path.join(directory, d))])):
class_indices[cls] = i
cls_dir = os.path.join(directory, cls)
for root, _, files in os.walk(cls_dir):
for file in files:
if file.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp', '.gif')):
@@ -62,20 +60,14 @@ def get_file_list(directory):
labels.append(i)
return file_paths, labels, class_indices
def calculate_class_weights(labels):
counts = Counter(labels)
total = len(labels)
num_classes = len(counts)
weights = {cls: total / (num_classes * count) for cls, count in counts.items()}
return weights
####################################################################################################
def multi_conv(x, strides):
convs = []
for k in [(1, 7), (7, 1), (3, 3), (5, 5)]:
for k in [(1, 7), (1, 1), (3, 3), (5, 5)]:
c = layers.DepthwiseConv2D(kernel_size=k, strides=strides, padding='same', activation='gelu')(x)
convs.append(c)
x = layers.Concatenate()(convs)
x = layers.BatchNormalization()(x)
return layers.Activation('gelu')(x)
@@ -90,17 +82,21 @@ def coordinate_attention(x, reduction):
h, w = input_shape[1], input_shape[2]
filters = x.shape[-1]
mip = max(8, filters // reduction)
x_h = layers.Lambda(lambda t: tf.reduce_mean(t, axis=2, keepdims=True))(x)
x_w = layers.Lambda(lambda t: tf.reduce_mean(t, axis=1, keepdims=True))(x)
x_w_t = layers.Permute((2, 1, 3))(x_w)
x_cat = layers.Concatenate(axis=1)([x_h, x_w_t])
x_cat = layers.Conv2D(mip, 1, padding='valid', use_bias=False)(x_cat)
x_cat = layers.BatchNormalization()(x_cat)
x_cat = layers.Activation('gelu')(x_cat)
x_cat_h, x_cat_w = layers.Lambda(lambda t: tf.split(t[0], [t[1], t[2]], axis=1))([x_cat, h, w])
x_h_sigmoid = layers.Conv2D(filters, 1, padding='valid', activation='sigmoid')(x_cat_h)
x_cat_w = layers.Permute((2, 1, 3))(x_cat_w)
x_w_sigmoid = layers.Conv2D(filters, 1, padding='valid', activation='sigmoid')(x_cat_w)
return layers.Multiply()([x, x_h_sigmoid, x_w_sigmoid])
def build_model(input_shape, num_classes):
@@ -111,8 +107,11 @@ def build_model(input_shape, num_classes):
x = multi_conv(x, strides=2)
x = multi_conv(x, strides=2)
x = multi_conv(x, strides=2)
x = coordinate_attention(x, reduction=4)
x = layers.GlobalMaxPooling2D()(x)
x = layers.Dense(x.shape[-1] // 2, activation='gelu', dtype='float32')(x)
o = layers.Dense(num_classes, activation='softmax', dtype='float32')(x)
return models.Model(i, o)
@@ -125,6 +124,14 @@ def calculate_lr_cosine_sq(epoch, total_epochs=25, lr_start=5e-4, lr_end=1e-6):
def lr_scheduler(epoch):
return calculate_lr_cosine_sq(epoch)
def categorical_focal_loss(gamma=2.0, alpha=0.25):
def focal_loss(y_true, y_pred):
y_pred = tf.clip_by_value(y_pred, tf.keras.backend.epsilon(), 1.0 - tf.keras.backend.epsilon())
cross_entropy = -y_true * tf.math.log(y_pred)
loss = alpha * tf.math.pow(1.0 - y_pred, gamma) * cross_entropy
return tf.reduce_sum(loss, axis=-1)
return focal_loss
class CustomModelCheckpoint(callbacks.Callback):
def on_epoch_end(self, epoch, logs=False):
models.save_model(self.model, f"PYAS_Model_Epoch_{epoch + 1}.h5")
@@ -140,19 +147,17 @@ for imgs, _ in train_ds.take(1):
train_ds.image_shape = imgs.shape[1:]
break
strategy = tf.distribute.MirroredStrategy(cross_device_ops=tf.distribute.HierarchicalCopyAllReduce())
with strategy.scope():
try:
model = models.load_model('PYAS_Model_Epoch_0.h5')
print("Load model from disk")
except:
print("Creating a new model")
model = build_model(train_ds.image_shape, len(train_ds.class_indices))
print(f"Total parameters: {model.count_params()}")
model.compile(optimizer='adam', loss='categorical_crossentropy',
metrics=['accuracy', metrics.Precision(name='precision'), metrics.Recall(name='recall')])
try:
model = models.load_model('PYAS_Model_Epoch_0.h5')
print("Load model from disk")
except:
print("Creating a new model")
model = build_model(train_ds.image_shape, len(train_ds.class_indices))
model.fit(train_ds, epochs=25, callbacks=[CustomModelCheckpoint(),
callbacks.LearningRateScheduler(lr_scheduler)], validation_data=val_ds)
print(f"Total parameters: {model.count_params()}")
model.compile(optimizer='adam', loss=categorical_focal_loss(gamma=2.0, alpha=0.25),
metrics=['accuracy', metrics.Precision(name='precision'), metrics.Recall(name='recall'),
metrics.FalsePositives(name='fp'), metrics.TrueNegatives(name='tn')])
model.fit(train_ds, epochs=25, callbacks=[CustomModelCheckpoint(), callbacks.LearningRateScheduler(lr_scheduler)], validation_data=val_ds)
input('Training Complete')