##两种初始化方法,一种利用apply,一种利用.modules 进行遍历#------------------apply---------------------------# 1. 根据网络层的不同定义不同的初始化方式 defweight_init(m):ifisinstance(m, nn.Linear):
nn.init.xavier_normal_(m.weight)
nn.init.constant_(m.bias,0)# 也可以判断是否为conv2d,使用相应的初始化方式 elifisinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')# 是否为批归一化层elifisinstance(m, nn.BatchNorm2d):
nn.init.constant_(m.weight,1)
nn.init.constant_(m.bias,0)# 2. 初始化网络结构
model = Net(in_dim, n_hidden_1, n_hidden_2, out_dim)# 3. 将weight_init应用在子模块上
model.apply(weight_init)#torch中的apply函数通过可以不断遍历model的各个模块。实际上其使用的是深度优先算法 #-------------.modules 进行遍历-------------------defkaiming_normal_init_weight(model):for m in model.modules():ifisinstance(m, nn.Conv2d):
torch.nn.init.kaiming_normal_(m.weight)elifisinstance(m, nn.BatchNorm2d):
m.weight.data.fill_(1)
m.bias.data.zero_()return model
节省显存
#节省显存操作#1.-----------删除变量------------for batch_idx,(images, targets)inenumerate(dataloader):
optimizer.zero_grad()
loss = model(images, targets)
loss.backward()
optimizer.step()del loss #zheli!!#forward时候减少使用变量defforward(self, x):
x = self.conv_1(x)
x = self.conv_2(x)return x
#2.-----------inplace操作-----------
nn.ReLU(inplace=True)#尽量在原tensor所在的内存空间上修改tensor,而不是copy一份进行操作#3.---------验证操作时关闭梯度--------with torch.no_grad():# 使用model进行预测的代码pass
import torch.nn.functional as F
import torch.nn as nn
import torch
from einops import rearrange
classLoss(nn.Module):def__init__(self, n_classes, weight=None, alpha=0.5):"dice_loss_plus_cetr_weighted"super(Loss, self).__init__()
self.n_classes = n_classes
self.weight = weight.cuda()# self.weight = weight
self.alpha = alpha
defforward(self,input, target):
smooth =0.01# 防止分母为0
input1 = F.softmax(input, dim=1)
target1 = F.one_hot(target,self.n_classes)
input1 = rearrange(input1,'b n h w s -> b n (h w s)')
target1 = rearrange(target1,'b h w s n -> b n (h w s)')
input1 = input1[:,1:,:]
target1 = target1[:,1:,:].float()# 以batch为单位计算loss和dice_loss,据说训练更稳定,和上面的公式有出入# 注意,这里的dice不是真正的dice,叫做soft_dice更贴切
inter = torch.sum(input1 * target1)
union = torch.sum(input1)+ torch.sum(target1)+ smooth
dice =2.0* inter / union
loss = F.cross_entropy(input,target, weight=self.weight)
total_loss =(1- self.alpha)* loss +(1- dice)* self.alpha
return total_loss
if __name__ =='__main__':
torch.manual_seed(3)
device = torch.device('cuda'if torch.cuda.is_available()else'cpu')
losser = Loss(n_classes=4, weight=torch.tensor([0.2,0.3,0.25,0.25])).to(device)
x = torch.randn((2,4,16,16,16)).to(device)
y = torch.randint(0,4,(2,16,16,16)).to(device)print(losser(x, y))
分类之label smooth
classCELoss(nn.Module):''' Cross Entropy Loss with label smoothing '''def__init__(self, label_smooth=None, class_num=10):super().__init__()
self.label_smooth = label_smooth
self.class_num = class_num
defforward(self, pred, target):'''
Args:
pred: prediction of model output [N, M]
target: ground truth of sampler [N, C]
'''
eps =1e-12if self.label_smooth isnotNone:# cross entropy loss with label smoothing
logprobs = F.log_softmax(pred, dim=1)# softmax + log# target = F.one_hot(target, self.class_num) # 转换成one-hot# label smoothing# 实现 1# target = (1.0-self.label_smooth)*target + self.label_smooth/self.class_num# 实现 2# implement 2
target = torch.clamp(target.float(),min=self.label_smooth /(self.class_num -1),max=1.0- self.label_smooth)
loss =-1* torch.sum(target * logprobs,1)else:# standard cross e ntropy loss
loss =-1.* pred.gather(1, target.unsqueeze(-1))+ torch.log(torch.exp(pred + eps).sum(dim=1))return loss.mean()
网络保存
defmakedirs(dirs):ifnot os.path.exists(dirs):
os.makedirs(dirs)# Path to save models
checkpoint_dir = os.path.join(
CONFIG.EXP.OUTPUT_DIR,"models",
CONFIG.EXP.ID,
CONFIG.MODEL.NAME.lower(),
CONFIG.DATASET.SPLIT.TRAIN,)# Save a modelif iteration % CONFIG.SOLVER.ITER_SAVE ==0:
torch.save(
model.module.state_dict(),
os.path.join(checkpoint_dir,"checkpoint_{}.pth".format(iteration)),)withopen("./socre.txt","w")as f:
json.dump(score, f, indent=4, sort_keys=True)#其中score是类似的字典形式return{"Pixel Accuracy": acc,"Mean Accuracy": acc_cls,"Frequency Weighted IoU": fwavacc,"Mean IoU": mean_iu,"Class IoU": cls_iu,}