diff --git a/Engine/Model/train_torch.py b/Engine/Model/train_torch.py index 19a576f..a9f7170 100644 --- a/Engine/Model/train_torch.py +++ b/Engine/Model/train_torch.py @@ -73,7 +73,7 @@ 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)] + kernel_sizes = [(1, 7), (7, 1), (1, 1), (3, 3), (5, 5)] for h, w in kernel_sizes: pad_h = (h - 1) // 2 pad_w = (w - 1) // 2 @@ -182,17 +182,17 @@ class PYASModel(nn.Module): 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.layer1 = ResidualBlock(c_in, 32, strides=1) + self.layer2 = ResidualBlock(32, 64, strides=2) + self.layer3 = ResidualBlock(64, 128, strides=2) + self.layer4 = ResidualBlock(128, 256, strides=2) + self.layer5 = ResidualBlock(256, 512, strides=2) + self.layer6 = ResidualBlock(512, 1024, strides=2) - self.se = SEBlock(512, reduction=4) + self.se = SEBlock(1024, 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)) + self.head = nn.Sequential(nn.Linear(2048, 1024), nn.GELU(), nn.Linear(1024, num_classes)) def forward(self, x): x = self.layer1(x) @@ -221,24 +221,32 @@ class ONNXExportWrapper(nn.Module): #################################################################################################### class CategoricalFocalLoss(nn.Module): - def __init__(self, alpha, gamma=2.0): + def __init__(self, alpha, gamma=2.0, smoothing=0.1): super().__init__() self.register_buffer('alpha', torch.tensor(alpha)) self.gamma = gamma + self.smoothing = smoothing 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() + num_classes = inputs.size(-1) + log_pt = F.log_softmax(inputs, dim=-1) + + with torch.no_grad(): + true_dist = torch.zeros_like(log_pt) + true_dist.fill_(self.smoothing / (num_classes - 1)) + true_dist.scatter_(1, targets.data.unsqueeze(1), 1.0 - self.smoothing) + + pt = torch.exp(log_pt) + focal_weight = self.alpha[targets].unsqueeze(1) * (1 - pt) ** self.gamma + loss = (-true_dist * log_pt * focal_weight).sum(dim=-1) + return loss.mean() #################################################################################################### def train(): data_dir = r'.\Image_File_Pefile' image_size = (224, 224) - batch_size = 128 + batch_size = 64 val_split = 0.00001 total_epochs = 30 lr_start = 1e-3