u52f85420
04bf90c09b
- 新增 irstd 核心包、scripts、tests、configs - 新增研究文档：literature_review / mathematical_analysis / novelty_matrix / experiment_plan / research_decision / research_log / manuscript_outline - 新增 evidence、results 数据与审计产物 - 新增 pyproject.toml、requirements、.gitignore
52 lines
2.1 KiB
Python
52 lines
2.1 KiB
Python
"""Original compact baseline, NOT a reproduction of MSHNet or DNANet."""
|
|
import torch
|
|
from torch import nn
|
|
from torch.nn import functional as F
|
|
|
|
class TinyUNet(nn.Module):
|
|
def __init__(self, width=16):
|
|
super().__init__()
|
|
def block(a,b):
|
|
return nn.Sequential(nn.Conv2d(a,b,3,padding=1),nn.ReLU(),
|
|
nn.Conv2d(b,b,3,padding=1),nn.ReLU())
|
|
self.enc = block(1,width)
|
|
self.mid = block(width,width*2)
|
|
self.dec = block(width*3,width)
|
|
self.head = nn.Conv2d(width,1,1)
|
|
def forward(self,x):
|
|
a = self.enc(x)
|
|
b = self.mid(F.max_pool2d(a,2,ceil_mode=True))
|
|
b = F.interpolate(b,size=a.shape[-2:],mode='bilinear',align_corners=False)
|
|
return self.head(self.dec(torch.cat((a,b),1)))
|
|
|
|
class CandidateGate(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.net=nn.Sequential(nn.Linear(5,16),nn.ReLU(),nn.Linear(16,1))
|
|
def forward(self,x):
|
|
return self.net(x).squeeze(-1)
|
|
|
|
def segmentation_loss(logits, target, name='bce_dice'):
|
|
p = logits.sigmoid()
|
|
bce = F.binary_cross_entropy_with_logits(logits,target)
|
|
dims=tuple(range(1,p.ndim)); eps=1e-6
|
|
inter=(p*target).sum(dims)
|
|
if name=='bce': return bce
|
|
if name=='focal':
|
|
pt=p*target+(1-p)*(1-target)
|
|
weight=.25*target+.75*(1-target)
|
|
return (weight*(1-pt).pow(2)*F.binary_cross_entropy_with_logits(logits,target,reduction='none')).mean()
|
|
dice=(1-(2*inter+eps)/(p.sum(dims)+target.sum(dims)+eps)).mean()
|
|
if name=='dice': return dice
|
|
if name=='bce_dice': return bce+dice
|
|
if name=='iou': return (1-(inter+eps)/(p.sum(dims)+target.sum(dims)-inter+eps)).mean()
|
|
raise ValueError(name)
|
|
|
|
def gate_loss(logits, labels, areas, dual, alpha, size_weight=True, tau=1.):
|
|
import math
|
|
pos=labels>0.5; neg=~pos
|
|
weight=areas.clamp_min(1).rsqrt().clamp(.1,1) if size_weight else torch.ones_like(areas)
|
|
positive=(F.softplus(-logits[pos])*weight[pos]).sum()/max(1,int(pos.sum()))
|
|
upper=F.softplus(logits[neg]/tau).sum()/math.log(2)
|
|
return positive+dual*(upper-alpha), upper
|