基于pytorch-lightning+monai的深度学习模型参数初始化

import os
import random
import numpy as np
import torch
import monai
import pytorch_lightning as pl

from monai.utils import set_determinism
from timm.models.layers import trunc_normal_

random.seed(42)
np.random.seed(42)
torch.manual_seed(42)
pl.seed_everything(42)
set_determinism(42)

def ModelParamsInit(model):
    assert isinstance(model, nn.Module)
    for m in model.modules():
        if isinstance(m, (nn.Conv3d, nn.ConvTranspose3d, nn.Linear)):
            trunc_normal_(m.weight, std=0.02)
            if m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, (nn.LayerNorm, nn.BatchNorm3d)):
            if m.weight is not None:
                nn.init.constant_(m.weight, 1.0)
            if m.bias is not None:
                nn.init.constant_(m.bias, 0)

使用这个模型参数初始化，可以很大程度上加速模型的训练速度。本人手动实现了一个模型，不加参数初始化时，模型在0-70个epoch之间的dice系数一直为0；加了参数初始化之后，第2个epoch结束之后dice系数达到了0.6左右。