if opt.test_dec_epoch != 0:
# Load pretrained models
decision_net.load_state_dict(
torch.load("./saved_models/decision_net_%d.pth" %
(opt.test_dec_epoch)))
transforms_ = transforms.Compose([
transforms.Resize((opt.img_height, opt.img_width), Image.BICUBIC),
transforms.ToTensor(),
#transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
])
testloader = DataLoader(
KolektorDataset(dataSetRoot,
transforms_=transforms_,
transforms_mask=None,
subFold="Test",
isTrain=False),
batch_size=1,
shuffle=False,
num_workers=0,
)
#segment_net.eval()
#decision_net.eval()
save_path_str = "testResult"
if os.path.exists(save_path_str) == False:
os.makedirs(save_path_str, exist_ok=True)
for i, testBatch in enumerate(testloader):
transforms_ = transforms.Compose([
transforms.Resize((opt.img_height, opt.img_width), Image.BICUBIC),
transforms.ToTensor(),
#transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
])
transforms_mask = transforms.Compose([
transforms.Resize((opt.img_height // 8, opt.img_width // 8)),
transforms.ToTensor(),
#transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
])
trainOKloader = DataLoader(
KolektorDataset(dataSetRoot,
transforms_=transforms_,
transforms_mask=transforms_mask,
subFold="Train_OK",
isTrain=True),
batch_size=opt.batch_size,
shuffle=True,
num_workers=opt.worker_num,
)
trainNGloader = DataLoader(
KolektorDataset(dataSetRoot,
transforms_=transforms_,
transforms_mask=transforms_mask,
subFold="Train_NG",
isTrain=True),
batch_size=opt.batch_size,
shuffle=True,
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--cuda",
type=bool,
default=True,
help="number of gpu")
parser.add_argument("--gpu_num", type=int, default=1, help="number of gpu")
parser.add_argument("--worker_num",
type=int,
default=4,
help="number of input workers")
parser.add_argument("--batch_size",
type=int,
default=4,
help="batch size of input")
parser.add_argument("--lr",
type=float,
default=0.0001,
help="adam: learning rate")
parser.add_argument("--b1",
type=float,
default=0.5,
help="adam: decay of first order momentum of gradient")
parser.add_argument("--b2",
type=float,
default=0.999,
help="adam: decay of first order momentum of gradient")
parser.add_argument("--begin_epoch",
type=int,
default=0,
help="begin_epoch")
parser.add_argument("--end_epoch", type=int, default=61, help="end_epoch")
parser.add_argument("--seg_epoch",
type=int,
default=100,
help="pretrained segment epoch")
parser.add_argument("--need_test",
type=bool,
default=True,
help="need to test")
parser.add_argument("--test_interval",
type=int,
default=10,
help="interval of test")
parser.add_argument("--need_save",
type=bool,
default=True,
help="need to save")
parser.add_argument("--save_interval",
type=int,
default=10,
help="interval of save weights")
parser.add_argument("--img_height",
type=int,
default=1408,
help="size of image height") # 1408x512 704x256
parser.add_argument("--img_width",
type=int,
default=512,
help="size of image width")
opt = parser.parse_args()
print(opt)
dataSetRoot = "./KolektorSDD_Data"
# Build nets
segment_net = SegmentNet(init_weights=True)
decision_net = DecisionNet(init_weights=True)
# Loss functions
#criterion_segment = torch.nn.MSELoss()
criterion_decision = torch.nn.MSELoss()
#criterion_decision = torch.nn.BCEWithLogitsLoss()
# Optimizers
optimizer_dec = torch.optim.Adam(decision_net.parameters(),
lr=opt.lr,
betas=(opt.b1, opt.b2))
if opt.cuda:
segment_net = segment_net.cuda()
decision_net = decision_net.cuda()
# criterion_segment.cuda()
criterion_decision.cuda()
if opt.gpu_num > 1:
segment_net = torch.nn.DataParallel(segment_net,
device_ids=list(range(
opt.gpu_num)))
decision_net = torch.nn.DataParallel(decision_net,
device_ids=list(range(
opt.gpu_num)))
if opt.begin_epoch != 0:
# Load pretrained models
decision_net.load_state_dict(
torch.load("./saved_models/decision_net_%d.pth" %
(opt.begin_epoch)))
else:
# Initialize weights
decision_net.apply(weights_init_normal)
# load pretrained segment parameters
segment_net.load_state_dict(
torch.load("./saved_models/segment_net_%d.pth" % (opt.seg_epoch)))
segment_net.eval()
transforms_ = transforms.Compose([
transforms.Resize((opt.img_height, opt.img_width), Image.BICUBIC),
transforms.ToTensor(),
#transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
])
transforms_mask = transforms.Compose([
transforms.Resize((opt.img_height // 8, opt.img_width // 8)),
transforms.ToTensor(),
#transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
])
trainOKloader = DataLoader(
KolektorDataset(dataSetRoot,
transforms_=transforms_,
transforms_mask=transforms_mask,
subFold="Train_OK",
isTrain=True),
batch_size=opt.batch_size,
shuffle=True,
num_workers=opt.worker_num,
)
trainNGloader = DataLoader(
KolektorDataset(dataSetRoot,
transforms_=transforms_,
transforms_mask=transforms_mask,
subFold="Train_NG",
isTrain=True),
batch_size=opt.batch_size,
shuffle=True,
num_workers=opt.worker_num,
)
'''
trainloader = DataLoader(
KolektorDataset(dataSetRoot, transforms_=transforms_, transforms_mask= transforms_mask,
subFold="Train_ALL", isTrain=True),
batch_size=opt.batch_size,
shuffle=True,
num_workers=opt.worker_num,
)
'''
testloader = DataLoader(
KolektorDataset(dataSetRoot,
transforms_=transforms_,
transforms_mask=transforms_mask,
subFold="Test/Train_NG",
isTrain=False),
batch_size=1,
shuffle=False,
num_workers=0,
)
for epoch in range(opt.begin_epoch, opt.end_epoch):
iterOK = trainOKloader.__iter__()
iterNG = trainNGloader.__iter__()
lenNum = min(len(trainNGloader), len(trainOKloader))
lenNum = 2 * (lenNum - 1)
decision_net.train()
segment_net.eval()
# train
for i in range(0, lenNum):
if i % 2 == 0:
batchData = iterOK.__next__()
#idx, batchData = enumerate(trainOKloader)
gt_c = Variable(torch.Tensor(
np.zeros((batchData["img"].size(0), 1))),
requires_grad=False)
else:
batchData = iterNG.__next__()
gt_c = Variable(torch.Tensor(
np.ones((batchData["img"].size(0), 1))),
requires_grad=False)
#idx, batchData = enumerate(trainNGloader)
if opt.cuda:
img = batchData["img"].cuda()
mask = batchData["mask"].cuda()
gt_c = gt_c.cuda()
else:
img = batchData["img"]
mask = batchData["mask"]
rst = segment_net(img)
f = rst["f"]
seg = rst["seg"]
optimizer_dec.zero_grad()
rst_d = decision_net(f, seg)
# rst_d = torch.Tensor.long(rst_d)
loss_dec = criterion_decision(rst_d, gt_c)
loss_dec.backward()
optimizer_dec.step()
sys.stdout.write(
"\r [Epoch %d/%d] [Batch %d/%d] [loss %f]" %
(epoch, opt.end_epoch, i, lenNum, loss_dec.item()))
# test
if opt.need_test and epoch % opt.test_interval == 0 and epoch >= opt.test_interval:
decision_net.eval()
# segment_net.eval()
all_time = 0
for i, testBatch in enumerate(testloader):
imgTest = testBatch["img"].cuda()
t1 = time.time()
rstTest = segment_net(imgTest)
fTest = rstTest["f"]
segTest = rstTest["seg"]
cTest = decision_net(fTest, segTest)
t2 = time.time()
save_path_str = "./testResultDec/epoch_%d" % epoch
if os.path.exists(save_path_str) == False:
os.makedirs(save_path_str, exist_ok=True)
# os.mkdir(save_path_str)
if cTest.item() > 0.5:
labelStr = "NG"
else:
labelStr = "OK"
save_image(imgTest.data,
"%s/img_%d_%s.jpg" % (save_path_str, i, labelStr))
save_image(
segTest.data,
"%s/img_%d_seg_%s.jpg" % (save_path_str, i, labelStr))
# print("processing image NO %d, time comsuption %fs"%(i, t2 - t1))
all_time = (t2 - t1) + all_time
count_time = i + 1
# print(all_time, count_time)
avg_time = all_time / count_time
print("\na image avg time %fs" % avg_time)
decision_net.train()
# save model parameters
if opt.need_save and epoch % opt.save_interval == 0 and epoch >= opt.save_interval:
# decision_net.eval()
save_path_str = "./saved_models"
if os.path.exists(save_path_str) == False:
os.makedirs(save_path_str, exist_ok=True)
torch.save(decision_net.state_dict(),
"%s/decision_net_%d.pth" % (save_path_str, epoch))
print("save weights ! epoch = %d" % epoch)
# decision_net.train()
pass