|
|
|
|
@ -0,0 +1,527 @@
|
|
|
|
|
import os
|
|
|
|
|
import torch
|
|
|
|
|
import torchvision
|
|
|
|
|
from torch.autograd import Variable
|
|
|
|
|
import torch.nn as nn
|
|
|
|
|
import torch.nn.functional as F
|
|
|
|
|
|
|
|
|
|
from torch.utils.data import Dataset, DataLoader
|
|
|
|
|
from torchvision import transforms, utils
|
|
|
|
|
import torch.optim as optim
|
|
|
|
|
import torchvision.transforms as standard_transforms
|
|
|
|
|
|
|
|
|
|
import numpy as np
|
|
|
|
|
import glob
|
|
|
|
|
import os
|
|
|
|
|
|
|
|
|
|
from data_loader import Rescale
|
|
|
|
|
from data_loader import RescaleT
|
|
|
|
|
from data_loader import RandomCrop
|
|
|
|
|
from data_loader import ToTensor
|
|
|
|
|
from data_loader import ToTensorLab
|
|
|
|
|
from data_loader import SalObjDataset
|
|
|
|
|
|
|
|
|
|
from skimage import io, transform
|
|
|
|
|
from PIL import Image
|
|
|
|
|
|
|
|
|
|
from model import U2NET
|
|
|
|
|
from model import U2NETP
|
|
|
|
|
from model import U2NETM
|
|
|
|
|
import cv2
|
|
|
|
|
|
|
|
|
|
# ------- 1. define loss function --------
|
|
|
|
|
|
|
|
|
|
bce_loss = nn.BCELoss(size_average=True)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
# normalize the predicted SOD probability map
|
|
|
|
|
def normPRED(d):
|
|
|
|
|
ma = torch.max(d)
|
|
|
|
|
mi = torch.min(d)
|
|
|
|
|
|
|
|
|
|
dn = (d-mi)/(ma-mi)
|
|
|
|
|
|
|
|
|
|
return dn
|
|
|
|
|
|
|
|
|
|
def IOU(pred, label):
|
|
|
|
|
|
|
|
|
|
Iand1 = np.sum(label[:, :]*pred[:, :])
|
|
|
|
|
Ior1 = np.sum(label[:, :]) + np.sum(pred[:, :]) - Iand1
|
|
|
|
|
IoU1 = Iand1/Ior1
|
|
|
|
|
|
|
|
|
|
return IoU1
|
|
|
|
|
# temp = np.abs(label - pred)
|
|
|
|
|
|
|
|
|
|
# union_minus_inter = np.sum(temp > 5)
|
|
|
|
|
|
|
|
|
|
# union_plus_inter = h * w * 2
|
|
|
|
|
|
|
|
|
|
# union = ((union_minus_inter + union_plus_inter)/2) - minus
|
|
|
|
|
|
|
|
|
|
# inter = union - union_minus_inter - minus
|
|
|
|
|
|
|
|
|
|
# print("inter: %d, union: %d" % (inter, union))
|
|
|
|
|
|
|
|
|
|
# iou = inter / union
|
|
|
|
|
|
|
|
|
|
return iou
|
|
|
|
|
|
|
|
|
|
# -----------------------------------修改loss函数，增加canny loss--------------------------------------------------------
|
|
|
|
|
def muti_bce_loss_fusion(d0, d1, d2, d3, d4, d5, d6, labels_v, mask_canny, label_canny):
|
|
|
|
|
loss0 = bce_loss(d0, labels_v)
|
|
|
|
|
loss1 = bce_loss(d1, labels_v)
|
|
|
|
|
loss2 = bce_loss(d2, labels_v)
|
|
|
|
|
loss3 = bce_loss(d3, labels_v)
|
|
|
|
|
loss4 = bce_loss(d4, labels_v)
|
|
|
|
|
loss5 = bce_loss(d5, labels_v)
|
|
|
|
|
loss6 = bce_loss(d6, labels_v)
|
|
|
|
|
loss_canny = bce_loss(mask_canny, label_canny)
|
|
|
|
|
loss = loss0 + loss1 + loss2 + loss3 + loss4 + loss5 + loss6 + loss_canny
|
|
|
|
|
print("l_canny: %3f, loss0: %3f, loss1: %3f, loss2: %3f, loss3: %3f, loss4: %3f, loss5: %3f, loss6: %3f\n" % (
|
|
|
|
|
loss_canny.data.item(), loss0.data.item(), loss1.data.item(), loss2.data.item(), loss3.data.item(), loss4.data.item(),
|
|
|
|
|
loss5.data.item(),
|
|
|
|
|
loss6.data.item()))
|
|
|
|
|
return loss0, loss
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
# ------- 2. set the directory of training dataset --------
|
|
|
|
|
|
|
|
|
|
model_name = 'u2net' # or 'u2net' or 'u2netm' or 'u2netp'
|
|
|
|
|
|
|
|
|
|
data_dir = os.path.join(os.getcwd(), 'train_data' + os.sep)
|
|
|
|
|
tra_image_dir = os.path.join('DUTS', 'DUTS-TR', 'miou_im' + os.sep)
|
|
|
|
|
tra_label_dir = os.path.join('DUTS', 'DUTS-TR', 'miou_gt' + os.sep)
|
|
|
|
|
|
|
|
|
|
image_ext = '.jpg'
|
|
|
|
|
label_ext = '.png'
|
|
|
|
|
|
|
|
|
|
model_dir = os.path.join(os.getcwd(), 'saved_models', model_name + os.sep)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
epoch_num = 666
|
|
|
|
|
batch_size_train = 10
|
|
|
|
|
print("---batch size = ", batch_size_train)
|
|
|
|
|
batch_size_val = 1
|
|
|
|
|
train_num = 0
|
|
|
|
|
val_num = 0
|
|
|
|
|
|
|
|
|
|
tra_img_name_list = glob.glob(data_dir + tra_image_dir + '*' + image_ext)
|
|
|
|
|
|
|
|
|
|
tra_lbl_name_list = []
|
|
|
|
|
for img_path in tra_img_name_list:
|
|
|
|
|
img_name = img_path.split(os.sep)[-1]
|
|
|
|
|
aaa = img_name.split(".")
|
|
|
|
|
bbb = aaa[0:-1]
|
|
|
|
|
imidx = bbb[0]
|
|
|
|
|
for i in range(1, len(bbb)):
|
|
|
|
|
imidx = imidx + "." + bbb[i]
|
|
|
|
|
|
|
|
|
|
tra_lbl_name_list.append(data_dir + tra_label_dir + imidx + label_ext)
|
|
|
|
|
|
|
|
|
|
print("---")
|
|
|
|
|
print("train images: ", len(tra_img_name_list))
|
|
|
|
|
print("train labels: ", len(tra_lbl_name_list))
|
|
|
|
|
print("---")
|
|
|
|
|
|
|
|
|
|
train_num = len(tra_img_name_list)
|
|
|
|
|
|
|
|
|
|
salobj_dataset = SalObjDataset(
|
|
|
|
|
img_name_list=tra_img_name_list,
|
|
|
|
|
lbl_name_list=tra_lbl_name_list,
|
|
|
|
|
transform=transforms.Compose([
|
|
|
|
|
RescaleT(320),
|
|
|
|
|
RandomCrop(288),
|
|
|
|
|
ToTensorLab(flag=0)]))
|
|
|
|
|
# compose 对列表中的操作依次执行，最后返回处理后的img
|
|
|
|
|
salobj_dataloader = DataLoader(salobj_dataset, batch_size=batch_size_train, shuffle=True, num_workers=0)
|
|
|
|
|
|
|
|
|
|
# ------- 3. define model --------
|
|
|
|
|
# define the net
|
|
|
|
|
if (model_name == 'u2net'):
|
|
|
|
|
net = U2NET(3, 1)
|
|
|
|
|
elif (model_name == 'u2netp'):
|
|
|
|
|
net = U2NETP(3, 1)
|
|
|
|
|
elif (model_name == 'u2netm'):
|
|
|
|
|
net = U2NETM(3, 1)
|
|
|
|
|
|
|
|
|
|
if torch.cuda.is_available():
|
|
|
|
|
net.cuda()
|
|
|
|
|
|
|
|
|
|
learning_rate = 1.e-7
|
|
|
|
|
# ------- 4. define optimizer --------
|
|
|
|
|
print("---defining optimizer...")
|
|
|
|
|
optimizer = optim.Adam(net.parameters(), lr=learning_rate, betas=(0.9, 0.999), eps=1e-08, weight_decay=0)
|
|
|
|
|
print('---optimizer defined successfully---')
|
|
|
|
|
|
|
|
|
|
# ------- 5. training process --------
|
|
|
|
|
print("---start TRAINING...")
|
|
|
|
|
|
|
|
|
|
# ------- 5(1). load model data --------
|
|
|
|
|
|
|
|
|
|
# c = input('? Do you need to LOAD the former model data ? (y/n): ')
|
|
|
|
|
# if c == 'y':
|
|
|
|
|
# # 加载模型参数
|
|
|
|
|
if torch.cuda.is_available():
|
|
|
|
|
# net.load_state_dict(torch.load(model_dir + model_name + '.pth'))
|
|
|
|
|
net = torch.load(model_dir + model_name + '.pth')
|
|
|
|
|
print('---GPU::the former model data LOADED successfully---')
|
|
|
|
|
else:
|
|
|
|
|
net.load_state_dict(torch.load(model_dir + model_name + '.pth', map_location='cpu'))
|
|
|
|
|
print('---CPU::the former model data LOADED successfully---')
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
# elif c == 'n':
|
|
|
|
|
# print('---the former model data IGNORED successfully---')
|
|
|
|
|
# else:
|
|
|
|
|
# print('ERROR ! invalid command !!!')
|
|
|
|
|
# print('---continue to TRAIN the model\nGood LUCK!...')
|
|
|
|
|
# ------------------------------------------
|
|
|
|
|
|
|
|
|
|
ite_num = 0
|
|
|
|
|
running_loss = 0.0
|
|
|
|
|
running_tar_loss = 0.0
|
|
|
|
|
ite_num4val = 0
|
|
|
|
|
save_frq = 10 # save the model every 10 iterations
|
|
|
|
|
iou_sum = 0
|
|
|
|
|
n = 0.
|
|
|
|
|
|
|
|
|
|
for epoch in range(0, epoch_num):
|
|
|
|
|
net.train()
|
|
|
|
|
|
|
|
|
|
for i, data in enumerate(salobj_dataloader):
|
|
|
|
|
n = n + 1.
|
|
|
|
|
ite_num = ite_num + 1
|
|
|
|
|
ite_num4val = ite_num4val + 1
|
|
|
|
|
|
|
|
|
|
if ite_num % 250 == 0:
|
|
|
|
|
learning_rate = learning_rate / 2.
|
|
|
|
|
print("---Optimizer: learning_rate REDUCED by half successfully---")
|
|
|
|
|
print('---NOTE: present learning_rate = ', learning_rate)
|
|
|
|
|
print("---GOOD LUCK---")
|
|
|
|
|
else:
|
|
|
|
|
learning_rate = learning_rate
|
|
|
|
|
optimizer = optim.Adam(net.parameters(), lr=learning_rate, betas=(0.9, 0.999), eps=1e-08, weight_decay=0)
|
|
|
|
|
|
|
|
|
|
inputs, labels = data['image'], data['label']
|
|
|
|
|
|
|
|
|
|
inputs = inputs.type(torch.FloatTensor)
|
|
|
|
|
print(inputs.size())
|
|
|
|
|
labels = labels.type(torch.FloatTensor)
|
|
|
|
|
|
|
|
|
|
# wrap them in Variable
|
|
|
|
|
if torch.cuda.is_available():
|
|
|
|
|
inputs_v, labels_v = Variable(inputs.cuda(), requires_grad=False), Variable(labels.cuda(), requires_grad=False)
|
|
|
|
|
else:
|
|
|
|
|
inputs_v, labels_v = Variable(inputs, requires_grad=False), Variable(labels, requires_grad=False)
|
|
|
|
|
|
|
|
|
|
# y zero the parameter gradients
|
|
|
|
|
optimizer.zero_grad()
|
|
|
|
|
|
|
|
|
|
# forward + backward + optimize
|
|
|
|
|
d0, d1, d2, d3, d4, d5, d6 = net(inputs_v)
|
|
|
|
|
# ------compute mask_canny and label_canny------------------------------------------------------------------------------
|
|
|
|
|
mask_c = d0[:, 0, :, :]
|
|
|
|
|
mask_c = normPRED(mask_c)
|
|
|
|
|
mask_c = mask_c.squeeze()
|
|
|
|
|
mask_c = mask_c.cpu().data.numpy()
|
|
|
|
|
mask_c = Image.fromarray(np.uint8(mask_c[0] * 255)) # PIL格式
|
|
|
|
|
canny_image_name = tra_img_name_list[i] # 原图地址
|
|
|
|
|
canny_image = io.imread(canny_image_name)
|
|
|
|
|
mask_c = mask_c.resize((canny_image.shape[1], canny_image.shape[0]), resample=Image.BILINEAR) # 改为原图尺寸
|
|
|
|
|
mask_c = np.asarray(mask_c)
|
|
|
|
|
mask_canny = cv2.Canny(mask_c, 100, 200, apertureSize=5)
|
|
|
|
|
label_c = cv2.imread(tra_lbl_name_list[i])
|
|
|
|
|
label_canny = cv2.Canny(label_c, 100, 200, apertureSize=5)
|
|
|
|
|
# 对输入bce loss的mask canny和label canny转为list再转为tensor
|
|
|
|
|
# bce loss要求传入的参数为list tensor并且数据大小为0-1
|
|
|
|
|
mask_canny = np.array(mask_canny)/255.0
|
|
|
|
|
mask_canny = mask_canny.tolist()
|
|
|
|
|
mask_canny = torch.Tensor(mask_canny)
|
|
|
|
|
label_canny = np.array(label_canny)/255.0
|
|
|
|
|
label_canny = label_canny.tolist()
|
|
|
|
|
label_canny = torch.Tensor(label_canny)
|
|
|
|
|
loss2, loss = muti_bce_loss_fusion(d0, d1, d2, d3, d4, d5, d6, labels_v, mask_canny, label_canny)
|
|
|
|
|
# 记得修改最上面的loss函数
|
|
|
|
|
# ------canny loss end here---------------------------------------------------------------------------------------------
|
|
|
|
|
loss.backward()
|
|
|
|
|
optimizer.step()
|
|
|
|
|
|
|
|
|
|
# print statistics
|
|
|
|
|
running_loss += loss.data.item()
|
|
|
|
|
running_tar_loss += loss2.data.item()
|
|
|
|
|
|
|
|
|
|
# print(tra_lbl_name_list[i])
|
|
|
|
|
# pred = d0[0, 0, :, :]
|
|
|
|
|
# pred = normPRED(pred)
|
|
|
|
|
# iou_image_name = tra_img_name_list[i]
|
|
|
|
|
# def create_pred(image_name, pred):
|
|
|
|
|
# predict = pred
|
|
|
|
|
# predict = predict.squeeze()
|
|
|
|
|
# predict_np = predict.cpu().data.numpy()
|
|
|
|
|
# print("predice_np:", predict_np.shape)
|
|
|
|
|
# predict_np = predict_np*255
|
|
|
|
|
# # im = Image.fromarray(np.uint8(predict_np[0]*255))
|
|
|
|
|
# img_name = image_name.split(os.sep)[-1]
|
|
|
|
|
# img_name = os.path.join(os.getcwd(), "train_data", tra_image_dir, img_name)
|
|
|
|
|
# print(img_name)
|
|
|
|
|
# image = cv2.imread(img_name, cv2.IMREAD_GRAYSCALE)
|
|
|
|
|
# print("image_shape", image.shape)
|
|
|
|
|
# pred = cv2.resize(predict_np, (image.shape[1], image.shape[0]))
|
|
|
|
|
# # image = io.imread(image_name)
|
|
|
|
|
# # imo = im.resize((image.shape[1], image.shape[0]), resample=Image.BILINEAR)
|
|
|
|
|
|
|
|
|
|
# print("pred_shape:", pred.shape)
|
|
|
|
|
# h = pred.shape[0]
|
|
|
|
|
# w = pred.shape[1]
|
|
|
|
|
# print(h, w)
|
|
|
|
|
# #二值化处理
|
|
|
|
|
# for i in range(h):
|
|
|
|
|
# for j in range(w):
|
|
|
|
|
# if pred[i][j] <= 127:
|
|
|
|
|
# pred[i][j] = 0
|
|
|
|
|
# if pred[i][j] > 127:
|
|
|
|
|
# pred[i][j] = 1
|
|
|
|
|
# return pred
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
# pred = create_pred(iou_image_name, pred)
|
|
|
|
|
|
|
|
|
|
# cv2.imwrite("pred_img.jpg", pred)
|
|
|
|
|
|
|
|
|
|
# iou_label_name = tra_lbl_name_list[i]
|
|
|
|
|
# h = pred.shape[0]
|
|
|
|
|
# w = pred.shape[1]
|
|
|
|
|
# mask = cv2.imread(iou_label_name, cv2.IMREAD_GRAYSCALE)
|
|
|
|
|
# #二值化处理
|
|
|
|
|
# for i in range(h):
|
|
|
|
|
# for j in range(w):
|
|
|
|
|
# if mask[i][j] <= 127:
|
|
|
|
|
# mask[i][j] = 0
|
|
|
|
|
# if mask[i][j] > 127:
|
|
|
|
|
# mask[i][j] = 1
|
|
|
|
|
# cv2.imwrite("mask_img.jpg", mask)
|
|
|
|
|
# # print("mask.shape: ", mask.shape) (300, 400)
|
|
|
|
|
# # mask_img = Image.fromarray(mask)
|
|
|
|
|
# # pred_img = Image.fromarray(pred)
|
|
|
|
|
# # mask_img.save("mask_img.jpg")
|
|
|
|
|
# # pred_img.save("pred_img.png")
|
|
|
|
|
# iou = IOU(pred, mask)
|
|
|
|
|
# iou_sum = iou_sum + iou
|
|
|
|
|
# MIOU = iou_sum / n
|
|
|
|
|
# print("MIOU = ", MIOU)
|
|
|
|
|
print("[epoch: %3d/%3d, batch: %5d/%5d, ite: %d]train loss: %3f, tar: %3f" % (
|
|
|
|
|
epoch + 1, epoch_num, (i + 1) * batch_size_train, train_num, ite_num, running_loss / ite_num4val,
|
|
|
|
|
running_tar_loss / ite_num4val))
|
|
|
|
|
# del temporary outputs and loss
|
|
|
|
|
del d0, d1, d2, d3, d4, d5, d6, loss2, loss
|
|
|
|
|
|
|
|
|
|
if ite_num % save_frq == 0:
|
|
|
|
|
torch.save(net, model_dir + model_name + '.pth')
|
|
|
|
|
# torch.save(model, f"../output/ckpt_%d.pth" % epoch)
|
|
|
|
|
# torch.save(net.state_dict(), model_dir + model_name + '.pth')
|
|
|
|
|
running_loss = 0.0
|
|
|
|
|
running_tar_loss = 0.0
|
|
|
|
|
net.train() # resume train
|
|
|
|
|
ite_num4val = 0
|
|
|
|
|
print('Congrats: model saved successfully!')
|
|
|
|
|
|
|
|
|
|
# ---------------------------------------------------------------------------------------------------------------------------------------------------------------------
|
|
|
|
|
# import os
|
|
|
|
|
# import torch
|
|
|
|
|
# import torchvision
|
|
|
|
|
# from torch.autograd import Variable
|
|
|
|
|
# import torch.nn as nn
|
|
|
|
|
# import torch.nn.functional as F
|
|
|
|
|
|
|
|
|
|
# from torch.utils.data import Dataset, DataLoader
|
|
|
|
|
# from torchvision import transforms, utils
|
|
|
|
|
# import torch.optim as optim
|
|
|
|
|
# import torchvision.transforms as standard_transforms
|
|
|
|
|
|
|
|
|
|
# import numpy as np
|
|
|
|
|
# import glob
|
|
|
|
|
# import os
|
|
|
|
|
|
|
|
|
|
# from data_loader import Rescale
|
|
|
|
|
# from data_loader import RescaleT
|
|
|
|
|
# from data_loader import RandomCrop
|
|
|
|
|
# from data_loader import ToTensor
|
|
|
|
|
# from data_loader import ToTensorLab
|
|
|
|
|
# from data_loader import SalObjDataset
|
|
|
|
|
|
|
|
|
|
# from model import U2NET
|
|
|
|
|
# from model import U2NETP
|
|
|
|
|
# from model import U2NETM
|
|
|
|
|
# # print(os.getcwd()) D:\A02分割论文\U-2-Net\U-2-Net-master
|
|
|
|
|
# # ------- 1. define loss function --------
|
|
|
|
|
|
|
|
|
|
# bce_loss = nn.BCELoss(size_average=True)
|
|
|
|
|
|
|
|
|
|
# def muti_bce_loss_fusion(d0, d1, d2, d3, d4, d5, d6, labels_v):
|
|
|
|
|
|
|
|
|
|
# loss0 = bce_loss(d0,labels_v)
|
|
|
|
|
# loss1 = bce_loss(d1,labels_v)
|
|
|
|
|
# loss2 = bce_loss(d2,labels_v)
|
|
|
|
|
# loss3 = bce_loss(d3,labels_v)
|
|
|
|
|
# loss4 = bce_loss(d4,labels_v)
|
|
|
|
|
# loss5 = bce_loss(d5,labels_v)
|
|
|
|
|
# loss6 = bce_loss(d6,labels_v)
|
|
|
|
|
|
|
|
|
|
# loss = loss0 + loss1 + loss2 + loss3 + loss4 + loss5 + loss6
|
|
|
|
|
# print("l0: %3f, l1: %3f, l2: %3f, l3: %3f, l4: %3f, l5: %3f, l6: %3f\n" % (loss0.data.item(), loss1.data.item(), loss2.data.item(), loss3.data.item(), loss4.data.item(), loss5.data.item(), loss6.data.item()))
|
|
|
|
|
# return loss0, loss
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
# # ------- 2. set the directory of training dataset --------
|
|
|
|
|
|
|
|
|
|
# model_name = 'u2netm' # or 'u2net' or 'u2netm' or 'u2netp'
|
|
|
|
|
|
|
|
|
|
# data_dir = os.path.join(os.getcwd(), 'train_data' + os.sep)
|
|
|
|
|
# tra_image_dir = os.path.join('DUTS', 'DUTS-TR', 'im_aug' + os.sep) # DUTS\DUTS-TR\im_aug\
|
|
|
|
|
# tra_label_dir = os.path.join('DUTS', 'DUTS-TR', 'gt_aug' + os.sep)
|
|
|
|
|
# # print(tra_image_dir) DUTS\DUTS-TR\im_aug\
|
|
|
|
|
# # print(data_dir) D:\A02图像分割\U-2-Net\U-2-Net-master\train_data\
|
|
|
|
|
# image_ext = '.jpg'
|
|
|
|
|
# label_ext = '.png'
|
|
|
|
|
|
|
|
|
|
# model_dir = os.path.join(os.getcwd(), 'saved_models', model_name + os.sep)
|
|
|
|
|
# # print(model_dir) # model_dir D:\A02图像分割\U-2-Net\U-2-Net-master\saved_models\u2netm\
|
|
|
|
|
|
|
|
|
|
# epoch_num = 666
|
|
|
|
|
# batch_size_train = 15
|
|
|
|
|
# batch_size_val = 1
|
|
|
|
|
# train_num = 0
|
|
|
|
|
# val_num = 0
|
|
|
|
|
|
|
|
|
|
# tra_img_name_list = glob.glob(data_dir + tra_image_dir + '*' + image_ext)
|
|
|
|
|
# # print(data_dir + tra_image_dir)
|
|
|
|
|
# # print(tra_img_name_list)
|
|
|
|
|
# tra_lbl_name_list = []
|
|
|
|
|
# for img_path in tra_img_name_list:
|
|
|
|
|
# img_name = img_path.split(os.sep)[-1]
|
|
|
|
|
# aaa = img_name.split(".")
|
|
|
|
|
# bbb = aaa[0:-1]
|
|
|
|
|
# imidx = bbb[0]
|
|
|
|
|
# for i in range(1,len(bbb)):
|
|
|
|
|
# imidx = imidx + "." + bbb[i]
|
|
|
|
|
|
|
|
|
|
# tra_lbl_name_list.append(data_dir + tra_label_dir + imidx + label_ext)
|
|
|
|
|
|
|
|
|
|
# print("---")
|
|
|
|
|
# print("train images: ", len(tra_img_name_list))
|
|
|
|
|
# print("train labels: ", len(tra_lbl_name_list))
|
|
|
|
|
# print("---")
|
|
|
|
|
|
|
|
|
|
# train_num = len(tra_img_name_list)
|
|
|
|
|
|
|
|
|
|
# salobj_dataset = SalObjDataset(
|
|
|
|
|
# img_name_list=tra_img_name_list,
|
|
|
|
|
# lbl_name_list=tra_lbl_name_list,
|
|
|
|
|
# transform=transforms.Compose([
|
|
|
|
|
# RescaleT(320),
|
|
|
|
|
# RandomCrop(288),
|
|
|
|
|
# ToTensorLab(flag=0)]))
|
|
|
|
|
|
|
|
|
|
# salobj_dataloader = DataLoader(salobj_dataset, batch_size=batch_size_train, shuffle=True, num_workers=0)
|
|
|
|
|
|
|
|
|
|
# # ------- 3. define model --------
|
|
|
|
|
# # define the net
|
|
|
|
|
# if(model_name == 'u2net'):
|
|
|
|
|
# net = U2NET(3,1)
|
|
|
|
|
# elif(model_name == 'u2netp'):
|
|
|
|
|
# net = U2NETP(3,1)
|
|
|
|
|
# elif(model_name == 'u2netm'):
|
|
|
|
|
# net = U2NETM(3,1)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
# if torch.cuda.is_available():
|
|
|
|
|
# net.cuda()
|
|
|
|
|
|
|
|
|
|
# learning_rate = 3.e-5
|
|
|
|
|
# # ------- 4. define optimizer --------
|
|
|
|
|
# print("---defining optimizer...")
|
|
|
|
|
# optimizer = optim.Adam(net.parameters(), lr=learning_rate, betas=(0.9, 0.999), eps=1e-08, weight_decay=0)
|
|
|
|
|
# print('---optimizer defined successfully---')
|
|
|
|
|
|
|
|
|
|
# # ------- 5. training process --------
|
|
|
|
|
# print("---start TRAINING...")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
# # ------- 5(1). load model data --------
|
|
|
|
|
# c = input('? Do you need to LOAD the former model data ? (y/n): ')
|
|
|
|
|
# if c == 'y':
|
|
|
|
|
# # 加载模型参数
|
|
|
|
|
# net.load_state_dict(torch.load(model_dir + model_name + '.pth'))
|
|
|
|
|
# print('---the former model data LOADED successfully---')
|
|
|
|
|
# elif c == 'n':
|
|
|
|
|
# print('---the former model data IGNORED successfully---')
|
|
|
|
|
# else:
|
|
|
|
|
# print('ERROR ! invalid command !!!')
|
|
|
|
|
# print('---continue to TRAIN the model\nGood LUCK!...')
|
|
|
|
|
|
|
|
|
|
# ite_num = 0
|
|
|
|
|
# running_loss = 0.0
|
|
|
|
|
# running_tar_loss = 0.0
|
|
|
|
|
# ite_num4val = 0
|
|
|
|
|
# save_frq = 10 # save the model every 20 iterations
|
|
|
|
|
|
|
|
|
|
# for epoch in range(0, epoch_num):
|
|
|
|
|
# net.train()
|
|
|
|
|
|
|
|
|
|
# for i, data in enumerate(salobj_dataloader):
|
|
|
|
|
# ite_num = ite_num + 1
|
|
|
|
|
# ite_num4val = ite_num4val + 1
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
# if ite_num % 250 == 0:
|
|
|
|
|
# learning_rate = learning_rate/2.
|
|
|
|
|
# print("---Optimizer: learning_rate REDUCED by half successfully---")
|
|
|
|
|
# print('---NOTE: present learning_rate = ', learning_rate)
|
|
|
|
|
# print("---GOOD LUCK---")
|
|
|
|
|
# else:
|
|
|
|
|
# learning_rate = learning_rate
|
|
|
|
|
# optimizer = optim.Adam(net.parameters(), lr=learning_rate, betas=(0.9, 0.999), eps=1e-08, weight_decay=0)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
# inputs, labels = data['image'], data['label']
|
|
|
|
|
|
|
|
|
|
# inputs = inputs.type(torch.FloatTensor)
|
|
|
|
|
# labels = labels.type(torch.FloatTensor)
|
|
|
|
|
|
|
|
|
|
# # wrap them in Variable
|
|
|
|
|
# if torch.cuda.is_available():
|
|
|
|
|
# inputs_v, labels_v = Variable(inputs.cuda(), requires_grad=False), Variable(labels.cuda(), requires_grad=False)
|
|
|
|
|
# else:
|
|
|
|
|
# inputs_v, labels_v = Variable(inputs, requires_grad=False), Variable(labels, requires_grad=False)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
# # y zero the parameter gradients
|
|
|
|
|
# optimizer.zero_grad()
|
|
|
|
|
|
|
|
|
|
# # forward + backward + optimize
|
|
|
|
|
# d0, d1, d2, d3, d4, d5, d6 = net(inputs_v)
|
|
|
|
|
# loss2, loss = muti_bce_loss_fusion(d0, d1, d2, d3, d4, d5, d6, labels_v)
|
|
|
|
|
|
|
|
|
|
# loss.backward()
|
|
|
|
|
# optimizer.step()
|
|
|
|
|
|
|
|
|
|
# # # print statistics
|
|
|
|
|
# running_loss += loss.data.item()
|
|
|
|
|
# running_tar_loss += loss2.data.item()
|
|
|
|
|
|
|
|
|
|
# # del temporary outputs and loss
|
|
|
|
|
# del d0, d1, d2, d3, d4, d5, d6, loss2, loss
|
|
|
|
|
|
|
|
|
|
# print("[epoch: %3d/%3d, batch: %5d/%5d, ite: %d] train loss: %3f, tar: %3f " % (
|
|
|
|
|
# epoch + 1, epoch_num, (i + 1) * batch_size_train, train_num, ite_num, running_loss / ite_num4val, running_tar_loss / ite_num4val))
|
|
|
|
|
|
|
|
|
|
# if ite_num % save_frq == 0:
|
|
|
|
|
|
|
|
|
|
# torch.save(net.state_dict(), model_dir + model_name + '.pth')
|
|
|
|
|
# running_loss = 0.0
|
|
|
|
|
# running_tar_loss = 0.0
|
|
|
|
|
# net.train() # resume train
|
|
|
|
|
# ite_num4val = 0
|
|
|
|
|
# print('Congrats: model saved successfully!')
|
