426 lines
16 KiB
Python
426 lines
16 KiB
Python
import os, random
|
|
import numpy as np
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
import torch.optim as optim
|
|
import torch.onnx
|
|
|
|
from torch.utils.data import Dataset, DataLoader
|
|
from torchvision import transforms
|
|
from PIL import Image
|
|
from tqdm import tqdm
|
|
|
|
####################################################################################################
|
|
|
|
Image.MAX_IMAGE_PIXELS = None
|
|
|
|
####################################################################################################
|
|
|
|
def get_file_list(directory):
|
|
file_paths, labels, class_indices = [], [], {}
|
|
dirs = sorted([d for d in os.listdir(directory) if os.path.isdir(os.path.join(directory, d))])
|
|
for i, cls in enumerate(dirs):
|
|
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')):
|
|
file_paths.append(os.path.join(root, file))
|
|
labels.append(i)
|
|
return file_paths, labels, class_indices
|
|
|
|
####################################################################################################
|
|
|
|
class CustomImageDataset(Dataset):
|
|
def __init__(self, file_paths, labels, image_size=(224, 224), color_mode="grayscale"):
|
|
self.file_paths = file_paths
|
|
self.labels = labels
|
|
self.image_size = image_size
|
|
self.color_mode = color_mode
|
|
self.transform = transforms.Resize(image_size)
|
|
|
|
def __len__(self):
|
|
return len(self.file_paths)
|
|
|
|
def __getitem__(self, idx):
|
|
path = self.file_paths[idx]
|
|
label = self.labels[idx]
|
|
try:
|
|
img = Image.open(path)
|
|
if self.color_mode == "grayscale":
|
|
img = img.convert("L")
|
|
else:
|
|
img = img.convert("RGB")
|
|
|
|
img = self.transform(img)
|
|
img_np = np.array(img).astype(np.float32) / 255.0
|
|
img_t = torch.from_numpy(img_np)
|
|
if img_t.ndim == 2:
|
|
img_t = img_t.unsqueeze(0)
|
|
elif img_t.ndim == 3:
|
|
img_t = img_t.permute(2, 0, 1)
|
|
|
|
except Exception:
|
|
c = 1 if self.color_mode == "grayscale" else 3
|
|
img_t = torch.zeros((c, *self.image_size), dtype=torch.float32)
|
|
return img_t, label
|
|
|
|
####################################################################################################
|
|
|
|
class MultiConv(nn.Module):
|
|
def __init__(self, in_channels, out_channels, strides):
|
|
super().__init__()
|
|
self.convs = nn.ModuleList()
|
|
kernel_sizes = [(1, 7), (3, 7), (1, 1), (3, 3)]
|
|
for h, w in kernel_sizes:
|
|
pad_h = (h - 1) // 2
|
|
pad_w = (w - 1) // 2
|
|
|
|
self.convs.append(nn.Sequential(
|
|
nn.Conv2d(in_channels, in_channels, kernel_size=(h, w), stride=strides,
|
|
padding=(pad_h, pad_w), dilation=1, groups=in_channels, bias=False),
|
|
nn.GELU(),
|
|
nn.BatchNorm2d(in_channels, eps=1e-3, momentum=0.01)))
|
|
|
|
self.concat_conv = nn.Conv2d(in_channels * len(kernel_sizes), out_channels, kernel_size=1, padding=0, bias=True)
|
|
self.act = nn.GELU()
|
|
self.bn = nn.BatchNorm2d(out_channels, eps=1e-3, momentum=0.01)
|
|
|
|
def forward(self, x):
|
|
outs = [conv(x) for conv in self.convs]
|
|
x = torch.cat(outs, dim=1)
|
|
x = self.concat_conv(x)
|
|
x = self.act(x)
|
|
x = self.bn(x)
|
|
return x
|
|
|
|
class CoordinateAttention(nn.Module):
|
|
def __init__(self, in_channels, reduction=4):
|
|
super().__init__()
|
|
mip = max(8, in_channels // reduction)
|
|
self.conv1 = nn.Conv2d(in_channels, mip, kernel_size=1, bias=True)
|
|
self.bn1 = nn.BatchNorm2d(mip)
|
|
self.act = nn.GELU()
|
|
self.conv_h = nn.Conv2d(mip, in_channels, kernel_size=1, bias=True)
|
|
self.conv_w = nn.Conv2d(mip, in_channels, kernel_size=1, bias=True)
|
|
|
|
def forward(self, x):
|
|
identity = x
|
|
n, c, h, w = x.size()
|
|
x_h = torch.mean(x, dim=3, keepdim=True)
|
|
x_w = torch.mean(x, dim=2, keepdim=True)
|
|
|
|
x_w_t = x_w.permute(0, 1, 3, 2)
|
|
y = torch.cat([x_h, x_w_t], dim=2)
|
|
|
|
y = self.conv1(y)
|
|
y = self.bn1(y)
|
|
y = self.act(y)
|
|
|
|
x_h_prime, x_w_prime = torch.split(y, [h, w], dim=2)
|
|
|
|
x_w_prime = x_w_prime.permute(0, 1, 3, 2)
|
|
|
|
a_h = torch.sigmoid(self.conv_h(x_h_prime))
|
|
a_w = torch.sigmoid(self.conv_w(x_w_prime))
|
|
|
|
return identity * a_h * a_w
|
|
|
|
class SEBlock(nn.Module):
|
|
def __init__(self, in_channels, reduction=4):
|
|
super().__init__()
|
|
self.conv1 = nn.Conv2d(in_channels, in_channels // reduction, kernel_size=1, bias=True)
|
|
self.act1 = nn.GELU()
|
|
self.bn1 = nn.BatchNorm2d(in_channels // reduction, eps=1e-3, momentum=0.01)
|
|
self.conv2 = nn.Conv2d(in_channels // reduction, in_channels, kernel_size=1, bias=True)
|
|
self.act2 = nn.Sigmoid()
|
|
|
|
def forward(self, x):
|
|
y = self.conv1(x)
|
|
y = self.act1(y)
|
|
y = self.bn1(y)
|
|
y = self.conv2(y)
|
|
y = self.act2(y)
|
|
return x * y
|
|
|
|
class ResidualBlock(nn.Module):
|
|
def __init__(self, in_channels, out_channels, strides):
|
|
super().__init__()
|
|
self.use_shortcut = (strides > 1) or (in_channels != out_channels)
|
|
|
|
if self.use_shortcut:
|
|
self.shortcut_pool = nn.MaxPool2d(kernel_size=strides, stride=strides) if strides > 1 else nn.Identity()
|
|
if in_channels != out_channels:
|
|
self.shortcut_conv = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False)
|
|
self.shortcut_bn = nn.BatchNorm2d(out_channels)
|
|
else:
|
|
self.shortcut_conv = nn.Identity()
|
|
self.shortcut_bn = nn.Identity()
|
|
|
|
self.multi_conv = MultiConv(in_channels, out_channels, strides)
|
|
self.coord_attn = CoordinateAttention(out_channels, reduction=4)
|
|
self.bn = nn.BatchNorm2d(out_channels)
|
|
self.final_act = nn.GELU()
|
|
|
|
def forward(self, x):
|
|
shortcut = x
|
|
if self.use_shortcut:
|
|
shortcut = self.shortcut_pool(shortcut)
|
|
shortcut = self.shortcut_conv(shortcut)
|
|
shortcut = self.shortcut_bn(shortcut)
|
|
|
|
x = self.multi_conv(x)
|
|
x = self.coord_attn(x)
|
|
x = self.bn(x)
|
|
x = x + shortcut
|
|
return self.final_act(x)
|
|
|
|
class PYASModel(nn.Module):
|
|
def __init__(self, input_shape, num_classes):
|
|
super().__init__()
|
|
c_in = input_shape[-1]
|
|
|
|
self.layer1 = ResidualBlock(c_in, 16, strides=1)
|
|
self.layer2 = ResidualBlock(16, 32, strides=2)
|
|
self.layer3 = ResidualBlock(32, 64, strides=2)
|
|
self.layer4 = ResidualBlock(64, 128, strides=2)
|
|
self.layer5 = ResidualBlock(128, 256, strides=2)
|
|
self.layer6 = ResidualBlock(256, 512, strides=2)
|
|
|
|
self.se = SEBlock(512, reduction=4)
|
|
self.gap = nn.AdaptiveAvgPool2d(1)
|
|
self.gmp = nn.AdaptiveMaxPool2d(1)
|
|
self.head = nn.Sequential(nn.Linear(1024, 1024), nn.GELU(), nn.Linear(1024, num_classes))
|
|
|
|
def forward(self, x):
|
|
x = self.layer1(x)
|
|
x = self.layer2(x)
|
|
x = self.layer3(x)
|
|
x = self.layer4(x)
|
|
x = self.layer5(x)
|
|
x = self.layer6(x)
|
|
x = self.se(x)
|
|
gap = torch.flatten(self.gap(x), 1)
|
|
gmp = torch.flatten(self.gmp(x), 1)
|
|
x = torch.cat([gap, gmp], dim=1)
|
|
x = self.head(x)
|
|
return x
|
|
|
|
class ONNXExportWrapper(nn.Module):
|
|
def __init__(self, model):
|
|
super().__init__()
|
|
self.model = model
|
|
|
|
def forward(self, x):
|
|
x = x.permute(0, 3, 1, 2)
|
|
x = self.model(x)
|
|
return F.softmax(x, dim=1)
|
|
|
|
####################################################################################################
|
|
|
|
class CategoricalFocalLoss(nn.Module):
|
|
def __init__(self, alpha, gamma=2.0):
|
|
super().__init__()
|
|
self.register_buffer('alpha', torch.tensor(alpha))
|
|
self.gamma = gamma
|
|
|
|
def forward(self, inputs, targets):
|
|
ce_loss = F.cross_entropy(inputs, targets, reduction='none')
|
|
pt = torch.exp(-ce_loss)
|
|
alpha_t = self.alpha[targets]
|
|
focal_loss = alpha_t * (1 - pt) ** self.gamma * ce_loss
|
|
return focal_loss.mean()
|
|
|
|
####################################################################################################
|
|
|
|
def train():
|
|
data_dir = r'.\Image_File_Pefile'
|
|
image_size = (224, 224)
|
|
batch_size = 128
|
|
val_split = 0.00001
|
|
total_epochs = 30
|
|
lr_start = 1e-3
|
|
lr_end = 1e-6
|
|
power = 2.5
|
|
color_mode = "grayscale"
|
|
|
|
file_paths, labels, class_indices = get_file_list(data_dir)
|
|
num_classes = len(class_indices)
|
|
sample_counts = {}
|
|
for lbl in labels:
|
|
sample_counts[lbl] = sample_counts.get(lbl, 0) + 1
|
|
total_samples = sum(sample_counts.values())
|
|
class_weights_dict = {k: total_samples / (num_classes * v) for k, v in sample_counts.items()}
|
|
alpha_list = [class_weights_dict[i] for i in range(num_classes)]
|
|
|
|
combined = list(zip(file_paths, labels))
|
|
random.shuffle(combined)
|
|
file_paths[:], labels[:] = zip(*combined)
|
|
|
|
num_val = int(val_split * len(file_paths))
|
|
train_files, train_labels = file_paths[num_val:], labels[num_val:]
|
|
val_files, val_labels = file_paths[:num_val], labels[:num_val]
|
|
|
|
train_ds = CustomImageDataset(train_files, train_labels, image_size, color_mode)
|
|
val_ds = CustomImageDataset(val_files, val_labels, image_size, color_mode)
|
|
|
|
train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True, num_workers=4, pin_memory=True, persistent_workers=True)
|
|
val_loader = DataLoader(val_ds, batch_size=batch_size, shuffle=False, num_workers=4, pin_memory=True, persistent_workers=True)
|
|
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
input_channels = 1 if color_mode == "grayscale" else 3
|
|
|
|
model = PYASModel(input_shape=(224, 224, input_channels), num_classes=num_classes)
|
|
if os.path.exists('PYAS_Model_Epoch_0.pth'):
|
|
print("Load model from disk")
|
|
model.load_state_dict(torch.load('PYAS_Model_Epoch_0.pth'))
|
|
else:
|
|
print("Creating a new model")
|
|
model = model.to(device)
|
|
|
|
if torch.cuda.device_count() > 1:
|
|
model = nn.DataParallel(model)
|
|
print(f"Total parameters: {sum(p.numel() for p in model.parameters() if p.requires_grad)}")
|
|
|
|
optimizer = optim.AdamW(model.parameters(), lr=lr_start, weight_decay=1e-4, amsgrad=True)
|
|
criterion = CategoricalFocalLoss(alpha_list, gamma=2.0).to(device)
|
|
scaler = torch.amp.GradScaler('cuda')
|
|
|
|
for epoch in range(total_epochs):
|
|
model.train()
|
|
|
|
remain_ratio = max(0, 1 - (epoch / total_epochs))
|
|
current_lr = lr_end + (lr_start - lr_end) * (remain_ratio ** power)
|
|
for param_group in optimizer.param_groups:
|
|
param_group['lr'] = current_lr
|
|
|
|
running_loss = 0.0
|
|
correct = 0
|
|
total = 0
|
|
|
|
tp_sum = torch.zeros(num_classes, device=device)
|
|
fp_sum = torch.zeros(num_classes, device=device)
|
|
fn_sum = torch.zeros(num_classes, device=device)
|
|
|
|
pbar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{total_epochs}")
|
|
for images, targets in pbar:
|
|
images, targets = images.to(device), targets.to(device)
|
|
current_batch_size = targets.size(0)
|
|
|
|
optimizer.zero_grad()
|
|
with torch.amp.autocast('cuda'):
|
|
outputs = model(images)
|
|
loss = criterion(outputs, targets)
|
|
|
|
scaler.scale(loss).backward()
|
|
scaler.step(optimizer)
|
|
scaler.update()
|
|
|
|
running_loss += loss.item() * current_batch_size
|
|
_, predicted = outputs.max(1)
|
|
total += current_batch_size
|
|
correct += predicted.eq(targets).sum().item()
|
|
|
|
for i in range(num_classes):
|
|
tp_sum[i] += ((predicted == i) & (targets == i)).sum()
|
|
fp_sum[i] += ((predicted == i) & (targets != i)).sum()
|
|
fn_sum[i] += ((predicted != i) & (targets == i)).sum()
|
|
|
|
avg_loss = running_loss / total
|
|
acc = correct / total
|
|
|
|
prec_cls = tp_sum / (tp_sum + fp_sum + 1e-7)
|
|
rec_cls = tp_sum / (tp_sum + fn_sum + 1e-7)
|
|
f1_cls = 2 * (prec_cls * rec_cls) / (prec_cls + rec_cls + 1e-7)
|
|
|
|
macro_prec = prec_cls.mean().item()
|
|
macro_f1 = f1_cls.mean().item()
|
|
|
|
mem = torch.cuda.memory_reserved(device) / (1024**3) if torch.cuda.is_available() else 0
|
|
|
|
pbar.set_postfix({
|
|
'loss': f'{avg_loss:.8f}',
|
|
'acc': f'{acc:.8f}',
|
|
'prec': f'{macro_prec:.8f}',
|
|
'f1': f'{macro_f1:.8f}',
|
|
'mem': f'{mem:.2f}G',
|
|
'lr': f'{current_lr:.8f}'
|
|
})
|
|
|
|
model.eval()
|
|
val_loss = 0.0
|
|
val_correct = 0
|
|
val_total = 0
|
|
v_tp = torch.zeros(num_classes, device=device)
|
|
v_fp = torch.zeros(num_classes, device=device)
|
|
v_fn = torch.zeros(num_classes, device=device)
|
|
|
|
val_pbar = tqdm(val_loader, desc=f"Val Epoch {epoch+1}")
|
|
|
|
with torch.no_grad():
|
|
for images, targets in val_pbar:
|
|
images, targets = images.to(device), targets.to(device)
|
|
current_batch_size = targets.size(0)
|
|
outputs = model(images)
|
|
loss = criterion(outputs, targets)
|
|
|
|
val_loss += loss.item() * current_batch_size
|
|
_, predicted = outputs.max(1)
|
|
val_total += current_batch_size
|
|
val_correct += predicted.eq(targets).sum().item()
|
|
|
|
for i in range(num_classes):
|
|
v_tp[i] += ((predicted == i) & (targets == i)).sum()
|
|
v_fp[i] += ((predicted == i) & (targets != i)).sum()
|
|
v_fn[i] += ((predicted != i) & (targets == i)).sum()
|
|
|
|
avg_val_loss = val_loss / val_total
|
|
val_acc = val_correct / val_total
|
|
|
|
v_prec_cls = v_tp / (v_tp + v_fp + 1e-7)
|
|
v_rec_cls = v_tp / (v_tp + v_fn + 1e-7)
|
|
v_f1_cls = 2 * (v_prec_cls * v_rec_cls) / (v_prec_cls + v_rec_cls + 1e-7)
|
|
|
|
val_macro_prec = v_prec_cls.mean().item()
|
|
val_macro_f1 = v_f1_cls.mean().item()
|
|
|
|
v_mem = torch.cuda.memory_reserved(device) / (1024**3) if torch.cuda.is_available() else 0
|
|
|
|
val_pbar.set_postfix({
|
|
'v_loss': f'{avg_val_loss:.8f}',
|
|
'v_acc': f'{val_acc:.8f}',
|
|
'v_prec': f'{val_macro_prec:.8f}',
|
|
'v_f1': f'{val_macro_f1:.8f}',
|
|
'v_mem': f'{v_mem:.2f}G',
|
|
'v_lr': f'{current_lr:.8f}'
|
|
})
|
|
|
|
model_to_save = model.module if isinstance(model, nn.DataParallel) else model
|
|
torch.save(model_to_save.state_dict(), f"PYAS_Model_Epoch_{epoch + 1}.pth")
|
|
|
|
onnx_wrapper = ONNXExportWrapper(model_to_save)
|
|
dummy_input = torch.randn(1, 224, 224, input_channels).to(device)
|
|
|
|
try:
|
|
torch.onnx.export(
|
|
onnx_wrapper,
|
|
dummy_input,
|
|
f"PYAS_Model_Epoch_{epoch + 1}.onnx",
|
|
export_params=True,
|
|
opset_version=17,
|
|
do_constant_folding=True,
|
|
input_names=['input'],
|
|
output_names=['output'],
|
|
dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}})
|
|
except Exception as e:
|
|
print(f"ONNX Export failed: {e}")
|
|
|
|
input('Training Complete')
|
|
|
|
####################################################################################################
|
|
|
|
if __name__ == '__main__':
|
|
train()
|