Add files via upload
This commit is contained in:
87owo authored and GitHub committed
1 parent
78323c9a08
commit
4163991b66
1 file changed
+425
@@ -0,0 +1,425 @@
|
||||
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()
|
||||
Reference in new issue
Block a user