else:
value.requires_grad = True
# for name, value in net.named_parameters():
# print('name: {0}, grad: {1}'.format(name, value.requires_grad))
net.load_state_dict(torch.load('./lib/vgg16_new1.model')) # 加载预训练好的模型
# net.load_state_dict(model_zoo.load_url(model_urls['vgg16']))
lib.utils.init_weight(net)
torch.nn.init.normal_(net.dd.weight, mean=0, std=0.01)
torch.nn.init.constant_(net.dd.bias, val=0)
if using_cuda:
net.cuda()
net.train() #启用 BatchNormalization 和 Dropout
print(net)
criterion = Loss.CTPN_Loss(using_cuda=using_cuda) # 获取loss
train_im_list, val_im_list = create_train_val() # 获取所有训练、测试数据文件地址列表
total_iter = len(train_im_list)
print("total training image num is %s" % len(train_im_list))
print("total val image num is %s" % len(val_im_list))
train_loss_list = []
test_loss_list = []
clc_loss_list = []
v_loss_list = []
side_loss_list = []
#########################################
################ 网络训练 ################
#########################################
# optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9, weight_decay=0.0005)
'cnn.VGG_16.convolution2_1.weight',
'cnn.VGG_16.convolution2_2.bias'
]
net = Net.CTPN()
for name, value in net.named_parameters():
if name in no_grad:
value.requires_grad = False
else:
value.requires_grad = True
net.load_state_dict(torch.load('./lib/vgg16.model'))
lib.utils.init_weight(net)
net.train()
print(net)
criterion = Loss.CTPN_Loss()
train_im_list, train_gt_list, val_im_list, val_gt_list = create_train_val()
total_iter = len(train_im_list)
print("total training image num is %s" % len(train_im_list))
print("total val image num is %s" % len(val_im_list))
train_loss_list = []
test_loss_list = []
for i in range(epoch):
if i >= change_epoch:
lr = lr_behind
else:
lr = lr_front
optimizer = optim.SGD(net.parameters(), lr=lr, momentum=0.9, weight_decay=0.0005)
iteration = 1