# 会占用大量显存
#---------------------#
aux_branch = False
#---------------------#
# 下采样的倍数
# 8和16
#---------------------#
downsample_factor = 16
#-------------------------------#
# Cuda的使用
#-------------------------------#
Cuda = True
model = PSPNet(num_classes=NUM_CLASSES,
backbone=backbone,
downsample_factor=downsample_factor,
pretrained=pretrained,
aux_branch=aux_branch).train()
# voc数据集下进行训练的
model_path = r"model_data/pspnet_mobilenetv2.pth"
# 加快模型训练的效率
print('Loading weights into state dict...')
model_dict = model.state_dict()
pretrained_dict = torch.load(model_path)
pretrained_dict = {
k: v
for k, v in pretrained_dict.items()
if np.shape(model_dict[k]) == np.shape(v)
}
model_dict.update(pretrained_dict)
#--------------------------------------------#
# 该部分代码只用于看网络结构，并非测试代码
#--------------------------------------------#
import torch
from torchsummary import summary
from nets.pspnet import PSPNet
if __name__ == "__main__":
model = PSPNet(num_classes=21,backbone="mobilenet",downsample_factor=16,aux_branch=False,pretrained=False).train().cuda()
summary(model,(3,473,473))
# 是否使用辅助分支
# 会占用大量显存
#------------------------------------------------------#
aux_branch = False
#------------------------------------------------------#
# 是否进行冻结训练，默认先冻结主干训练后解冻训练。
#------------------------------------------------------#
Freeze_Train = True
#------------------------------------------------------#
# 用于设置是否使用多线程读取数据
# 开启后会加快数据读取速度，但是会占用更多内存
# 内存较小的电脑可以设置为2或者0
#------------------------------------------------------#
num_workers = 4
model = PSPNet(num_classes=num_classes, backbone=backbone, downsample_factor=downsample_factor, pretrained=pretrained, aux_branch=aux_branch)
if not pretrained:
weights_init(model)
if model_path != '':
#------------------------------------------------------#
# 权值文件请看README，百度网盘下载
#------------------------------------------------------#
print('Load weights {}.'.format(model_path))
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model_dict = model.state_dict()
pretrained_dict = torch.load(model_path, map_location = device)
pretrained_dict = {k: v for k, v in pretrained_dict.items() if np.shape(model_dict[k]) == np.shape(v)}
model_dict.update(pretrained_dict)
model.load_state_dict(model_dict)
model_train = model.train()