深度学习代码技巧

深度学习代码技巧

网络初始化

##两种初始化方法，一种利用apply，一种利用.modules 进行遍历
#------------------apply---------------------------
# 1. 根据网络层的不同定义不同的初始化方式     
def weight_init(m):
    if isinstance(m, nn.Linear):
        nn.init.xavier_normal_(m.weight)
        nn.init.constant_(m.bias, 0)
    # 也可以判断是否为conv2d，使用相应的初始化方式 
    elif isinstance(m, nn.Conv2d):
        nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
     # 是否为批归一化层
    elif isinstance(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 进行遍历-------------------
def kaiming_normal_init_weight(model):
    for m in model.modules():
        if isinstance(m, nn.Conv2d):
            torch.nn.init.kaiming_normal_(m.weight)
        elif isinstance(m, nn.BatchNorm2d):
            m.weight.data.fill_(1)
            m.bias.data.zero_()
    return model

节省显存

#节省显存操作
#1.-----------删除变量------------
for batch_idx, (images, targets) in enumerate(dataloader):
    optimizer.zero_grad()
    loss = model(images, targets)
    loss.backward()
    optimizer.step()
    del loss #zheli!!
#forward时候减少使用变量
def forward(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

指定GPU

#-----------------指定GPU------------------
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1" 
#根据顺序表示优先使用0号设备,然后使用1号设备。

自定义学习率更新

def adjust_learning_rate(iter_num):
    lr = base_lr * (1.0 - iter_num / max_iterations) ** 0.9
    return lr


scheduler = torch.optim.lr_scheduler.StepLR(optimizer,step_size=30,gamma=0.1)
for epoch in range(epochs):
	for img, lab in train_loader:
		img,lab = img.cuda(), lab.cuda()
		pred = net(img)
		loss = criterion(pred, lab)
		#参数更新
		loss.backward()
		optimizer.step()
		if scheduler not None: # 注意这里后调用scheduler，避免第一个学习率改变
			scheduler.step() 
		elif update_way == 1 : # 通过更改optimizer的lr方式更新
			for param_group in optimizer.param_groups:
   				param_group["lr"] = adjust_learning_rate(iter)
   		else:
   			optimizer.param_groups[0]['lr'] = adjust_learning_rate(iter)

		#打印学习率
		print(optimizer.state_dict()['param_groups'][0]['lr'])

分割常用

one hot

#-----------------手动实现--------------------
import torch
class_num = 8
batch_size = 4

def one_hot(label):
    """
    将一维列表转换为独热编码
    """
    label = label.resize_(batch_size, 1)
    m_zeros = torch.zeros(batch_size, class_num)
    # 从 value 中取值，然后根据 dim 和 index 给相应位置赋值
    onehot = m_zeros.scatter_(1, label, 1)  # (dim,index,value)

    return onehot.numpy()  # Tensor -> Numpy

label = torch.LongTensor(batch_size).random_() % class_num  # 对随机数取余
print(one_hot(label))

# output:
[[0. 0. 0. 1. 0. 0. 0. 0.]
 [0. 0. 0. 0. 1. 0. 0. 0.]
 [0. 0. 1. 0. 0. 0. 0. 0.]
 [0. 1. 0. 0. 0. 0. 0. 0.]]

#-------------PyTorch自带------------------
import torch.nn.functional as F
import torch

tensor =  torch.arange(0, 5) % 3  # tensor([0, 1, 2, 0, 1])
one_hot = F.one_hot(tensor)

# 输出：
# tensor([[1, 0, 0],
#         [0, 1, 0],
#         [0, 0, 1],
#         [1, 0, 0],
#         [0, 1, 0]])

dice loss

import torch.nn.functional as F
import torch.nn as nn
import torch
from einops import rearrange


class Loss(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

    def forward(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

class CELoss(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

    def forward(self, pred, target):
        '''
        Args:
            pred: prediction of model output   [N, M]
            target: ground truth of sampler [N, C]
        '''
        eps = 1e-12
        if self.label_smooth is not None:
            # 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()

网络保存

def makedirs(dirs):
    if not 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 model
if iteration % CONFIG.SOLVER.ITER_SAVE == 0:
  torch.save(
    model.module.state_dict(),
    os.path.join(checkpoint_dir, "checkpoint_{}.pth".format(iteration)),
  )
  
with open("./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,
    }