def main(cfg, gpu):
torch.cuda.set_device(gpu)
# Network Builders
builder = ModelBuilder()
net_encoder = builder.build_encoder(arch=cfg.MODEL.arch_encoder.lower(),
fc_dim=cfg.MODEL.fc_dim,
weights=cfg.MODEL.weights_encoder)
net_decoder = builder.build_decoder(arch=cfg.MODEL.arch_decoder.lower(),
fc_dim=cfg.MODEL.fc_dim,
num_class=cfg.DATASET.num_class,
weights=cfg.MODEL.weights_decoder,
use_softmax=True)
crit = nn.NLLLoss(ignore_index=-1)
segmentation_module = SegmentationModule(net_encoder, net_decoder, crit)
# Dataset and Loader
dataset_val = ValDataset(cfg.DATASET.root_dataset, cfg.DATASET.list_val,
cfg.DATASET)
loader_val = torchdata.DataLoader(dataset_val,
batch_size=cfg.VAL.batch_size,
shuffle=False,
collate_fn=user_scattered_collate,
num_workers=5,
drop_last=True)
segmentation_module.cuda()
# Main loop
evaluate(segmentation_module, loader_val, cfg, gpu)
print('Evaluation Done!')
def worker(args, dev_id, start_idx, end_idx, result_queue):
torch.cuda.set_device(dev_id)
# Dataset and Loader
dataset_val = ValDataset(args.list_val,
args,
max_sample=args.num_val,
start_idx=start_idx,
end_idx=end_idx)
loader_val = torchdata.DataLoader(dataset_val,
batch_size=args.batch_size,
shuffle=False,
collate_fn=user_scattered_collate,
num_workers=2)
# Network Builders
builder = ModelBuilder()
net_encoder = builder.build_encoder(arch=args.arch_encoder,
fc_dim=args.fc_dim,
weights=args.weights_encoder)
net_decoder = builder.build_decoder(arch=args.arch_decoder,
fc_dim=args.fc_dim,
num_class=args.num_class,
weights=args.weights_decoder,
use_softmax=True)
crit = nn.NLLLoss(ignore_index=-1)
segmentation_module = SegmentationModule(net_encoder, net_decoder, crit)
segmentation_module.cuda()
# Main loop
evaluate(segmentation_module, loader_val, args, dev_id, result_queue)
def main(args):
torch.cuda.set_device(args.gpu_id)
# Network Builders
builder = ModelBuilder()
net_encoder = builder.build_encoder(arch=args.arch_encoder,
fc_dim=args.fc_dim,
weights=args.weights_encoder)
net_decoder = builder.build_decoder(arch=args.arch_decoder,
fc_dim=args.fc_dim,
weights=args.weights_decoder,
use_softmax=True)
crit = nn.NLLLoss(ignore_index=-1)
segmentation_module = SegmentationModule(net_encoder, net_decoder, crit)
# Dataset and Loader
dataset_val = ValDataset(args.list_val, args, max_sample=args.num_val)
loader_val = torchdata.DataLoader(dataset_val,
batch_size=args.batch_size,
shuffle=False,
collate_fn=user_scattered_collate,
num_workers=5,
drop_last=True)
segmentation_module.cuda()
# Main loop
evaluate(segmentation_module, loader_val, args)
print('Evaluation Done!')
def create_multi_source_train_data_loader(args):
training_records = broden_dataset.record_list['train']
# 0: object, part, scene
# 1: material
multi_source_iters = []
for idx_source in range(len(training_records)):
dataset = TrainDataset(training_records[idx_source], idx_source, args,
batch_per_gpu=args.batch_size_per_gpu)
loader_object_part_scene = torchdata.DataLoader(
dataset,
batch_size=args.num_gpus, # we have modified data_parallel
shuffle=False, # we do not use this param
collate_fn=user_scattered_collate,
num_workers=int(args.workers),
drop_last=True,
pin_memory=True)
multi_source_iters.append(iter(loader_object_part_scene))
# sample from multi source
nr_record = [len(records) for records in training_records]
sample_prob = np.asarray(nr_record) / np.sum(nr_record)
while True: # TODO(LYC):: set random seed.
source_idx = np.random.choice(len(training_records), 1, p=sample_prob)[0]
yield next(multi_source_iters[source_idx]), source_idx
def main(cfg):
#torch.cuda.set_device(gpu)
# Network Builders
builder = ModelBuilder()
net_encoder = builder.build_encoder(arch=cfg.MODEL.arch_encoder,
fc_dim=cfg.MODEL.fc_dim,
weights=cfg.MODEL.weights_encoder)
net_decoder = builder.build_decoder(arch=cfg.MODEL.arch_decoder,
fc_dim=cfg.MODEL.fc_dim,
num_class=cfg.DATASET.num_class,
weights=cfg.MODEL.weights_decoder,
use_softmax=True)
crit = nn.NLLLoss(ignore_index=-1)
segmentation_module = SegmentationModule(net_encoder, net_decoder, crit)
# Dataset and Loader
dataset_test = TestDataset(cfg.list_test, cfg.DATASET)
loader_test = torchdata.DataLoader(dataset_test,
batch_size=cfg.TEST.batch_size,
shuffle=False,
collate_fn=user_scattered_collate,
num_workers=5,
drop_last=True)
#segmentation_module.cuda()
# Main loop
test(segmentation_module, loader_test)
print('Inference done!')
def main(args):
torch.cuda.set_device(args.gpu_id)
# Network Builders
builder = ModelBuilder()
net_encoder = builder.build_encoder(arch=args.arch_encoder,
fc_dim=args.fc_dim,
weights=args.weights_encoder)
net_decoder = builder.build_decoder(arch=args.arch_decoder,
fc_dim=args.fc_dim,
nr_classes=args.nr_classes,
weights=args.weights_decoder,
use_softmax=True)
segmentation_module = SegmentationModule(net_encoder, net_decoder)
segmentation_module.cuda()
# Dataset and Loader
list_test = [{'fpath_img': args.test_img}]
dataset_val = TestDataset(list_test, args, max_sample=args.num_val)
loader_val = torchdata.DataLoader(dataset_val,
batch_size=args.batch_size,
shuffle=False,
collate_fn=user_scattered_collate,
num_workers=5,
drop_last=True)
# Main loop
test(segmentation_module, loader_val, args)
print('Inference done!')
def fetch_teacher_outputs(args):
teacher_outputs = []
# Network Builders
builder = ModelBuilder()
net_encoder = builder.build_encoder(arch=args.arch_encoder,
fc_dim=args.fc_dim,
weights=args.weights_encoder)
net_decoder = builder.build_decoder(arch=args.arch_decoder,
fc_dim=args.fc_dim,
num_class=args.num_class,
weights=args.weights_decoder)
crit = nn.NLLLoss(ignore_index=-1)
segmentation_module = SegmentationModule(net_encoder, net_decoder, crit)
# dataset and loader, use train
dataset_train = TrainDataset(args.list_train,
args,
batch_per_gpu=args.batch_size_per_gpu)
loader_train = torchdata.DataLoader(
dataset_train,
batch_size=args.num_gpus, # we have modified data_parallel
shuffle=False, # we do not use this param
collate_fn=user_scattered_collate,
num_workers=int(args.workers),
drop_last=True,
pin_memory=True)
segmentation_module.cuda()
# here may be some problems
for i, batch_data in enumerate(loader_train):
batch_data = batch_data[0] # get data list
seg_label = as_numpy(batch_data['seg_label'][0])
img_resized_list = batch_data['img_data']
with torch.no_grad():
segSize = (seg_label.shape[0], seg_label.shape[1])
# 预测输出的形状
pred = torch.zeros(1, args.num_class, segSize[0], segSize[1])
for img in img_resized_list:
feed_dict = batch_data.copy()
feed_dict['img_data'] = img
del feed_dict['img_ori']
del feed_dict['info']
feed_dict = async_copy_to(feed_dict, args.gpu_id)
# forward pass
pred_tmp = segmentation_module(feed_dict, segSize=segSize)
pred = pred + pred_tmp.cpu() / len(args.imgSize)
_, preds = torch.max(pred, dim=1)
preds.as_numpy(preds.squeeze(0))
teacher_outputs.append(preds)
return teacher_outputs
def main(args):
# Network Builders
is_finetune=True
builder = ModelBuilder()
encoder_fn=builder.build_encoder_textmountain
decoder_fn=builder.build_decoder_textmountain
dataset_train = TrainDataset(batch_per_gpu=12
)
loader_train = torchdata.DataLoader(
dataset_train,
batch_size=1, # we have modified data_parallel
shuffle=False, # we do not use this param
collate_fn=user_scattered_collate,
num_workers=int(2),
drop_last=True,
pin_memory=False)
iterator_train = iter(loader_train)
if is_finetune:
encoder_path='model_path'
decoder_path='model_path'
net_encoder = encoder_fn(encoder_path)
net_decoder = decoder_fn(decoder_path)
else:
net_encoder = encoder_fn(
)
net_decoder = decoder_fn()
segmentation_module=Module_textmountain(
net_encoder,net_decoder)
segmentation_module.cuda()
# Set up optimizers
nets = (net_encoder, net_decoder)
optimizers = create_optimizers(nets, args)
# Main loop
history = {'train': {'epoch': [], 'loss': [], 'acc': []}}
for epoch in range(args.start_epoch, args.num_epoch + 1):
train_epoch(segmentation_module, iterator_train, optimizers, history, epoch, args)
# checkpointing
checkpoint(optimizers,nets, history, args, epoch)
print('Training Done!')
def worker(args, dev_id, start_idx, end_idx, result_queue):
torch.cuda.set_device(dev_id)
# Dataset and Loader
dataset_val = ValDataset(args.list_val,
args,
max_sample=args.num_val,
start_idx=start_idx,
end_idx=end_idx)
loader_val = torchdata.DataLoader(dataset_val,
batch_size=args.batch_size,
shuffle=False,
collate_fn=user_scattered_collate,
num_workers=2)
# Network Builders
builder = ModelBuilder()
#net_encoder = builder.build_encoder(
# arch=args.arch_encoder,
# fc_dim=args.fc_dim,
# weights=args.weights_encoder)
#net_decoder = builder.build_decoder(
# arch=args.arch_decoder,
# fc_dim=args.fc_dim,
# num_class=args.num_class,
# weights=args.weights_decoder,
# use_softmax=True)
snet_encoder = builder.build_encoder(arch=cfg.SEM.ARCH_ENCODER,
fc_dim=cfg.SEM.FC_DIM)
net_decoder = builder.build_decoder(arch=cfg.SEM.DECODER_TYPE,
fc_dim=cfg.SEM.FC_DIM,
num_class=cfg.MODEL.NUM_CLASSES,
use_softmax=not self.training,
weights='')
crit = nn.NLLLoss(ignore_index=-1)
segmentation_module = SegmentationModule(net_encoder, net_decoder, crit)
print("Loading model weights")
state_dict = {}
pretrained = torch.load(args.ckpt,
map_location=lambda storage, loc: storage)
pretrained = pretrained['model']
segmentation_module.load_state_dict(pretrained, strict=True)
print("Weights load success")
segmentation_module.cuda()
# Main loop
evaluate(segmentation_module, loader_val, args, dev_id, result_queue)
def main(args):
# torch.cuda.set_device(args.gpu)
# Network Builders
builder = ModelBuilder()
enc_out = torch.randn(([1,2048,64,64]))
net_encoder = builder.build_encoder(
weights="baseline-resnet50dilated-ppm_deepsup/encoder_epoch_20.pth")
gcu = GraphConv()#, V=2), GCU(X=enc_out, V=4), GCU(X=enc_out, V=8),GCU(X=enc_out, V=32)]
crit = nn.NLLLoss(ignore_index=-1)
segmentation_module = SegmentationModule(net_encoder, gcu, crit, tr=False)
# print("Prinitng Params", gcu[1].parameters())
for m in gcu.parameters():
print("Hello",m.shape,m.name,m)
print("dddddddddddddddd", len(list(gcu.parameters())))
for m in gcu.modules():
print("Prining", m.parameters())
# Dataset and Loader
if len(args.test_imgs) == 1 and os.path.isdir(args.test_imgs[0]):
test_imgs = find_recursive(args.test_imgs[0])
else:
test_imgs = args.test_imgs
list_test = [{'fpath_img': x} for x in test_imgs]
dataset_test = TestDataset(
list_test, args, max_sample=-1)
loader_test = torchdata.DataLoader(
dataset_test,
batch_size=1,
shuffle=False,
collate_fn=user_scattered_collate,
num_workers=5,
drop_last=True)
# Main loop
test(segmentation_module, loader_test, args)
print('Inference done!')
def main(args):
torch.cuda.set_device(args.gpu)
# Network Builders
builder = ModelBuilder()
net_encoder = builder.build_encoder(
arch=args.arch_encoder,
fc_dim=args.fc_dim,
weights=args.weights_encoder)
net_decoder = builder.build_decoder(
arch=args.arch_decoder,
fc_dim=args.fc_dim,
num_class=args.num_class,
weights=args.weights_decoder,
use_softmax=True)
crit = nn.NLLLoss(ignore_index=-1)
segmentation_module = SegmentationModule(net_encoder, net_decoder, crit)
# Dataset and Loader
# list_test = [{'fpath_img': args.test_img}]
# test_chk = []
# testing = os.listdir("/home/teai/externalhd2/BDD100K/segmentation_v2/test/")
# for i in testing:
# if(i.endswith(".jpg")):
# test_chk.append("/home/teai/externalhd2/BDD100K/segmentation_v2/test/"+i)
test_chk = ['./05b07068-373666cb.jpg']
print(type(args.test_imgs))
list_test = [{'fpath_img': x} for x in test_chk]
#list_test=[{'fpath_img': 'frame_143.png'},{'fpath_img': 'frame_100.png'},{'fpath_img': 'frame_1.png'}]
#print("list_test",list_test)
dataset_test = TestDataset(
list_test, args, max_sample=args.num_val)
loader_test = torchdata.DataLoader(
dataset_test,
batch_size=args.batch_size,
shuffle=False,
collate_fn=user_scattered_collate,
num_workers=5,
drop_last=True)
segmentation_module.cuda()
# Main loop
start=time.time()
test(segmentation_module, loader_test, args)
end=time.time()
def main(args):
torch.cuda.set_device(args.gpu_id)
# Network Builders
builder = ModelBuilder()
net_encoder = builder.build_encoder(arch=args.arch_encoder,
fc_dim=args.fc_dim,
weights=args.weights_encoder)
net_decoder = builder.build_decoder(arch=args.arch_decoder,
fc_dim=args.fc_dim,
num_class=args.num_class,
weights=args.weights_decoder,
use_softmax=True)
crit = nn.NLLLoss(ignore_index=-1)
segmentation_module = SegmentationModule(net_encoder, net_decoder, crit)
# Dataset and Loader
# list_test = [{'fpath_img': args.test_img}]
# Pass a single image
if args.test_imgs is not None:
list_test = [{'fpath_img': x} for x in args.test_imgs]
if args.test_img_path is not None:
list_test = [{
'fpath_img': os.path.join(args.test_img_path, x)
} for x in os.listdir(args.test_img_path)]
list_test = list_test[:20]
dataset_test = TestDataset(list_test, args, max_sample=args.num_val)
loader_test = torchdata.DataLoader(dataset_test,
batch_size=args.batch_size,
shuffle=False,
collate_fn=user_scattered_collate,
num_workers=5,
drop_last=True)
segmentation_module.cuda()
# Main loop
test(segmentation_module, loader_test, args)
print('Inference done!')
def main(args):
torch.cuda.set_device(args.gpu)
# Network Builders
builder = ModelBuilder()
net_encoder = builder.build_encoder(
arch=args.arch_encoder,
fc_dim=args.fc_dim,
weights=args.weights_encoder)
net_decoder = builder.build_decoder(
arch=args.arch_decoder,
fc_dim=args.fc_dim,
num_class=args.num_class,
weights=args.weights_decoder,
use_softmax=True)
crit = nn.NLLLoss(ignore_index=-1)
segmentation_module = SegmentationModule(net_encoder, net_decoder, crit)
# Dataset and Loader
if len(args.test_imgs) == 1 and os.path.isdir(args.test_imgs[0]):
test_imgs = find_recursive(args.test_imgs[0])
else:
test_imgs = args.test_imgs
list_test = [{'fpath_img': x} for x in test_imgs]
dataset_test = TestDataset(
list_test, args, max_sample=args.num_val)
loader_test = torchdata.DataLoader(
dataset_test,
batch_size=args.batch_size,
shuffle=False,
collate_fn=user_scattered_collate,
num_workers=5,
drop_last=True)
segmentation_module.cuda()
# Main loop
test(segmentation_module, loader_test, args)
print('Inference done!')
def main(args):
torch.cuda.set_device(args.gpu)
crit = nn.NLLLoss(ignore_index=-1)
# Dataset and Loader
dataset_val = Dataset(args,
split_name=args.split_name,
batch_per_gpu=args.batchsize)
loader_val = torchdata.DataLoader(dataset_val,
batch_size=args.batchsize,
shuffle=False,
collate_fn=user_scattered_collate,
num_workers=5,
drop_last=True)
# Network Builders
builder = ModelBuilder()
print('Loading encoder from: %s' % (args.weights_encoder))
print('Loading decoder from: %s' % (args.weights_decoder))
net_encoder = builder.build_encoder(arch=args.arch_encoder,
fc_dim=args.fc_dim,
weights=args.weights_encoder)
net_decoder = builder.build_decoder(arch=args.arch_decoder,
fc_dim=args.fc_dim,
num_class=dataset_val.num_classes,
weights=args.weights_decoder,
use_softmax=True)
segmentation_module = SegmentationModule(net_encoder, net_decoder, crit)
segmentation_module.cuda()
# Main loop
if args.dataset.startswith('nyuv2sn'):
metrics = evaluate_surface_normals(segmentation_module, loader_val,
args)
else:
metrics = evaluate_segmentation(segmentation_module, loader_val, args)
save_pickle(args.result_file, metrics)
print('Evaluation Done!')
def main(args):
# Network Builders
builder = XceptionModelBuilder() ## NOTE: need to be changed!!! ##
net_encoder = builder.build_encoder(
arch=args.arch_encoder,
weights=args.weights_encoder,
overall_stride=args.overall_stride)
net_decoder = builder.build_decoder(
arch=args.arch_decoder,
fc_dim=2048,
num_class=args.num_class,
weights=args.weights_decoder)
crit = None
segmentation_module = SegmentationModule(net_encoder, net_decoder, crit)
print(segmentation_module)
# Dataset and Loader
dataset_test = VOCTestDataset(args)
loader_test = torchdata.DataLoader(
dataset_test,
batch_size=1, # data_parallel have been modified, not useful
shuffle=False, # do not use this param
collate_fn=user_scattered_collate,
num_workers=1, # MUST be 1 or 0
drop_last=False,
pin_memory=False)
# create loader iterator
iterator_test = iter(loader_test)
segmentation_module.cuda()
# Set up optimizers
nets = (net_encoder, net_decoder, crit)
test(segmentation_module, iterator_test, args)
print('[{}] Test Done!'.format(datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")))
def main(args):
# Network Builders
builder = XceptionModelBuilder()
net_encoder = builder.build_encoder(
arch=args.arch_encoder,
weights=args.weights_encoder,
overall_stride=args.overall_stride)
net_decoder = builder.build_decoder(
arch=args.arch_decoder,
fc_dim=2048,
num_class=args.num_class,
weights=args.weights_decoder)
crit = nn.CrossEntropyLoss(ignore_index=255, reduction='sum')
if args.arch_decoder.endswith('deepsup'):
segmentation_module = SegmentationModule(
net_encoder, net_decoder, crit, args.deep_sup_scale)
else:
segmentation_module = SegmentationModule(
net_encoder, net_decoder, crit)
print(segmentation_module)
segmentation_module.cuda()
# Set up optimizers
nets = (net_encoder, net_decoder, crit)
dataset_val = VOCValDataset(args) # create val dataset loader for every eval, in order to use drop_last=false
loader_val = torchdata.DataLoader(
dataset_val,
# data_parallel have been modified, MUST use val batch size
# and collate_fn MUST be user_scattered_collate
batch_size=args.val_batch_size,
shuffle=False,
collate_fn=user_scattered_collate,
num_workers=1, # MUST be 1 or 0
drop_last=False)
iterator_val = iter(loader_val)
(cls_ious, cls_mean_iou) = evaluate(segmentation_module, iterator_val, args)
print('[{}] Training Done!'.format(datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")))
def worker(cfg, gpu_id, start_idx, end_idx, result_queue):
#torch.cuda.set_device(gpu_id)
print("batch_size: " + str(cfg.VAL.batch_size))
# Dataset and Loader
dataset_val = ValDataset(cfg.DATASET.root_dataset,
cfg.DATASET.list_val,
cfg.DATASET,
start_idx=start_idx,
end_idx=end_idx)
loader_val = torchdata.DataLoader(dataset_val,
batch_size=cfg.VAL.batch_size,
shuffle=False,
collate_fn=user_scattered_collate,
num_workers=2)
# Network Builders
builder = ModelBuilder()
net_encoder = builder.build_encoder(arch=cfg.MODEL.arch_encoder.lower(),
fc_dim=cfg.MODEL.fc_dim,
weights=cfg.MODEL.weights_encoder)
net_decoder = builder.build_decoder(arch=cfg.MODEL.arch_decoder.lower(),
fc_dim=cfg.MODEL.fc_dim,
num_class=cfg.DATASET.num_class,
weights=cfg.MODEL.weights_decoder,
use_softmax=True)
crit = nn.NLLLoss(ignore_index=-1)
segmentation_module = SegmentationModule(net_encoder, net_decoder, crit)
#segmentation_module.cuda()
# Main loop
evaluate(segmentation_module, loader_val, cfg, gpu_id, result_queue)
def main():
suffix = '_epoch_25.pth'
test_imgs = ["input/n33.jpg"]
model_path = 'baseline-resnet101-upernet'
#model_path = 'a'
# Model related arguments
arch_encoder = 'resnet101'
arch_decoder = 'upernet'
#arch_encoder='resnet101dilated'
#arch_decoder='ppm_deepsup'
fc_dim = 2048
# Data related arguments
num_val = -1
num_class = 150
batch_size = 1
imgSize = [300, 400, 500, 600]
imgMaxSize = 1000
padding_constant = 8
segm_downsampling_rate = 8
# Misc arguments
result = '.',
gpu_id = 0
#torch.cuda.set_device(gpu_id)
# Network Builders
builder = ModelBuilder()
print('-----------1')
net_encoder = builder.build_encoder(arch=arch_encoder,
fc_dim=fc_dim,
weights=weights_encoder)
net_decoder = builder.build_decoder(arch=arch_decoder,
fc_dim=fc_dim,
num_class=num_class,
weights=weights_decoder,
use_softmax=True)
print('-----------------2')
crit = nn.NLLLoss(ignore_index=-1)
print('--------------------3')
segmentation_module = SegmentationModule(net_encoder, net_decoder, crit)
# Dataset and Loader
# list_test = [{'fpath_img': args.test_img}]
print('--------------------4')
list_test = [{'fpath_img': x} for x in test_imgs]
dataset_test = TestDataset(list_test, max_sample=num_val)
loader_test = torchdata.DataLoader(dataset_test,
batch_size=batch_size,
shuffle=False,
collate_fn=user_scattered_collate,
num_workers=5,
drop_last=True)
print('--------------------5')
segmentation_module
print('--------------------6')
# Main loop
pred = test(segmentation_module, loader_test)
print(pred)
print('Inference done!')
def main(args):
# Network Builders
builder = ModelBuilder()
net_encoder = builder.build_encoder(arch=args.arch_encoder,
fc_dim=args.fc_dim,
weights=args.weights_encoder)
net_decoder = builder.build_decoder(arch=args.arch_decoder,
fc_dim=args.fc_dim,
num_class=args.num_class,
weights=args.weights_decoder)
crit = nn.NLLLoss(ignore_index=-1)
if args.arch_decoder.endswith('deepsup'):
segmentation_module = SegmentationModule(net_encoder, net_decoder,
crit, args.deep_sup_scale)
else:
segmentation_module = SegmentationModule(net_encoder, net_decoder,
crit)
# Dataset and Loader
dataset_train = TrainDataset(args.list_train,
args,
batch_per_gpu=args.batch_size_per_gpu)
loader_train = torchdata.DataLoader(
dataset_train,
batch_size=args.num_gpus, # we have modified data_parallel
shuffle=False, # we do not use this param
collate_fn=user_scattered_collate,
num_workers=int(args.workers),
drop_last=True,
pin_memory=True)
print('1 Epoch = {} iters'.format(args.epoch_iters))
# create loader iterator
iterator_train = iter(loader_train)
# load nets into gpu
if args.num_gpus > 1:
segmentation_module = UserScatteredDataParallel(segmentation_module,
device_ids=range(
args.num_gpus))
# For sync bn
patch_replication_callback(segmentation_module)
segmentation_module.cuda()
# Set up optimizers
nets = (net_encoder, net_decoder, crit)
optimizers = create_optimizers(nets, args)
# Main loop
history = {'train': {'epoch': [], 'loss': [], 'acc': []}}
for epoch in range(args.start_epoch, args.num_epoch + 1):
train(segmentation_module, iterator_train, optimizers, history, epoch,
args)
# checkpointing
checkpoint(nets, history, args, epoch)
print('Training Done!')
def main(args):
torch.cuda.set_device(args.gpu)
# Network Builders
builder = ModelBuilder()
net_encoder = builder.build_encoder(arch=args.arch_encoder,
fc_dim=args.fc_dim,
weights=args.weights_encoder)
net_decoder = builder.build_decoder(arch=args.arch_decoder,
fc_dim=args.fc_dim,
num_class=args.num_class,
weights=args.weights_decoder,
use_softmax=True)
crit = nn.NLLLoss(ignore_index=-1)
segmentation_module = SegmentationModule(net_encoder, net_decoder, crit)
# Dataset and Loader
args.test_imgs = [
# 'test_input/panel_00_ori.jpg',
# 'test_input/panel_01_ori.jpg',
# 'test_input/panel_02_ori.jpg',
# 'test_input/panel_03_ori.jpg',
# 'test_input/panel_04_ori.jpg',
# 'test_input/panel_05_ori.jpg',
# 'test_input/panel_06_ori.jpg',
# 'test_input/panel_07_ori.jpg',
# 'test_input/panel_08_ori.jpg',
# 'test_input/panel_09_ori.jpg',
# 'test_input/panel_10_ori.jpg',
# 'test_input/panel_11_ori.jpg',
# 'test_input/panel_12_ori.jpg',
# 'test_input/panel_13_ori.jpg',
# 'test_input/173_768x1536.jpg',
# 'test_input/panel_81010.jpg',
# 'test_input/panel_81013.jpg',
# 'test_input/panel_81014.jpg',
# 'test_input/panel_81021.jpg',
# 'test_input/panel_81022.jpg',
# 'test_input/panel_81028.jpg',
# 'test_input/panel_81030.jpg',
'test_input/panel_81045.jpg',
'test_input/panel_81049.jpg',
'test_input/panel_81050.jpg',
'test_input/panel_81056.jpg',
'test_input/panel_81087.jpg',
'test_input/panel_81109.jpg',
'test_input/panel_81121.jpg',
'test_input/panel_81122.jpg',
'test_input/panel_81127.jpg',
'test_input/panel_81138.jpg',
'test_input/panel_81142.jpg',
'test_input/panel_81162.jpg',
'test_input/panel_81195.jpg',
'test_input/panel_81410.jpg',
'test_input/panel_84608.jpg'
]
list_test = [{'fpath_img': x} for x in args.test_imgs]
print(list_test)
dataset_test = TestDataset(list_test, args, max_sample=args.num_val)
loader_test = torchdata.DataLoader(dataset_test,
batch_size=args.batch_size,
shuffle=False,
collate_fn=user_scattered_collate,
num_workers=5,
drop_last=True)
segmentation_module.cuda()
# Main loop
test(segmentation_module, loader_test, args)
print('Inference done!')
def main():
# Network Builders
builder = ModelBuilder()
net_encoder = builder.build_encoder_textmountain
net_decoder = builder.build_decoder_textmountain
encoder_path = 'model_path'
decoder_path = 'model_path'
net_encoder = net_encoder(encoder_path)
net_decoder = net_decoder(decoder_path)
segmentation_module = Module_textmountain(net_encoder, net_decoder, ctw=1)
segmentation_module.cuda()
segmentation_module.eval()
dataset_train = ValDataset(size=800)
loader_train = torchdata.DataLoader(
dataset_train,
batch_size=1, # we have modified data_parallel
shuffle=False, # we do not use this param
collate_fn=user_scattered_collate,
num_workers=int(5),
drop_last=True,
pin_memory=True)
iterator_train = iter(loader_train)
for i_data in range(len(iterator_train)):
batch_data = next(iterator_train)
with torch.no_grad():
feed_dict = {}
feed_dict['images'] = batch_data[0]['im'].cuda()
pred_polygon = segmentation_module.forward(
feed_dict,
is_Train=False,
img_name=batch_data[0]['im_name'],
save_dir=save_dir)
if len(pred_polygon) == 0:
continue
im = batch_data[0]['im_name']
scale_w = batch_data[0]['scale_w']
scale_h = batch_data[0]['scale_h']
for polygon in pred_polygon:
polygon[:, 0] /= scale_w
polygon[:, 1] /= scale_h
if pred_polygon is not None:
res_file = os.path.join(
save_dir,
'{}.txt'.format(os.path.basename(im).split('.')[0]))
with open(res_file, 'w') as f:
for box in pred_polygon:
pred_txt = ''
box = box.reshape(-1)
for var in box:
pred_txt = pred_txt + '{},'.format(var)
pred_txt = pred_txt[:-1] + '\r\n'
f.write(pred_txt)
def main(args):
# Network Builders
builder = ModelBuilder()
net_encoder = builder.build_encoder(arch=args.arch_encoder,
fc_dim=args.fc_dim,
weights=args.weights_encoder)
net_decoder = builder.build_decoder(arch=args.arch_decoder,
fc_dim=args.fc_dim,
num_class=150,
weights=args.weights_decoder)
crit = nn.NLLLoss(ignore_index=-1)
if args.arch_decoder.endswith('deepsup'):
segmentation_module = SegmentationModule(net_encoder, net_decoder,
crit, args.deep_sup_scale)
else:
segmentation_module = SegmentationModule(net_encoder, net_decoder,
crit)
########
for param in segmentation_module.encoder.parameters():
param.requires_grad = False
#for name, param in segmentation_module.decoder.named_parameters():
# print(name)
# if(name == "conv_last.weight" or name =="conv_last.bias" or name =="conv_last_deepsup.weight" or name =="conv_last_deepsup.bias"):
# param.requires_grad = True
#else:
# param.requires_grad = False
#print(param.requires_grad)
segmentation_module.decoder.conv_last = nn.Conv2d(args.fc_dim // 4, 12, 1,
1, 0)
#segmentation_module.decoder.conv_last.
segmentation_module.decoder.conv_last_deepsup = nn.Conv2d(
args.fc_dim // 4, 12, 1, 1, 0)
########
# Dataset and Loader
dataset_train = TrainDataset(args.list_train,
args,
batch_per_gpu=args.batch_size_per_gpu)
loader_train = torchdata.DataLoader(
dataset_train,
batch_size=len(args.gpus), # we have modified data_parallel
shuffle=False, # we do not use this param
collate_fn=user_scattered_collate,
num_workers=int(args.workers),
drop_last=True,
pin_memory=True)
print('1 Epoch = {} iters'.format(args.epoch_iters))
# create loader iterator
iterator_train = iter(loader_train)
#######
#torch.backends.cudnn.benchmark = True
#CUDA_LAUNCH_BLOCKING=1
#######
# load nets into gpu
if len(args.gpus) > 1:
segmentation_module = UserScatteredDataParallel(segmentation_module,
device_ids=args.gpus)
# For sync bn
patch_replication_callback(segmentation_module)
segmentation_module.cuda()
# Set up optimizers
nets = (net_encoder, net_decoder, crit)
optimizers = create_optimizers(nets, args)
# Main loop
history = {'train': {'epoch': [], 'loss': [], 'acc': []}}
for epoch in range(args.start_epoch, args.num_epoch + 1):
train(segmentation_module, iterator_train, optimizers, history, epoch,
args)
# checkpointing
checkpoint(nets, history, args, epoch)
print('Training Done!')
def main(args):
# Network Builders
builder = ModelBuilder()
crit = nn.NLLLoss(ignore_index=-1)
crit = crit.cuda()
net_encoder = builder.build_encoder(
weights="baseline-resnet50dilated-ppm_deepsup/encoder_epoch_20.pth")
gcu = GraphConv(
batch=args.batch_size_per_gpu
) #, V=2), GCU(X=enc_out, V=4), GCU(X=enc_out, V=8),GCU(X=enc_out, V=32)]
# gcu.load_state_dict(torch.load("ckpt/baseline-resnet50dilated-ngpus1-batchSize1-imgMaxSize1000-paddingConst8-segmDownsampleRate8-epoch20/decoder_epoch_20.pth"))
segmentation_module = SegmentationModule(net_encoder, gcu, crit, tr=True)
# Dataset and Loader
dataset_train = TrainDataset(args.list_train,
args,
batch_per_gpu=args.batch_size_per_gpu)
loader_train = torchdata.DataLoader(
dataset_train,
batch_size=len(args.gpus), # we have modified data_parallel
shuffle=False, # we do not use this param
collate_fn=user_scattered_collate,
num_workers=int(args.workers),
drop_last=True,
pin_memory=True)
print('1 Epoch = {} iters'.format(args.epoch_iters))
# create loader iterator
iterator_train = iter(loader_train)
# load nets into gpu
if len(args.gpus) > 4:
segmentation_module = UserScatteredDataParallel(segmentation_module,
device_ids=args.gpus)
# For sync bn
patch_replication_callback(segmentation_module)
# segmentation_module.cuda()
# Set up optimizers
# print(gcu[0].parameters())
nets = (net_encoder, gcu, crit)
optimizers, par = create_optimizers(nets, args)
# Main loop
history = {'train': {'epoch': [], 'loss': [], 'acc': []}}
vis = visdom.Visdom()
win = vis.line(np.array([5.7]),
opts=dict(xlabel='epochs',
ylabel='Loss',
title='Training Loss V=16',
legend=['Loss']))
for epoch in range(args.start_epoch, args.num_epoch + 1):
lss = train(segmentation_module, iterator_train, optimizers, history,
epoch, par, vis, win, args)
# checkpointing
checkpoint(nets, history, args, epoch)
print('Training Done!')
self.train_cropSize = 321 # 512,328
self.random_flip = True
self.random_scale = True
self.random_scale_factor_max = 1.2
self.random_scale_factor_min = 0.75
self.train_list_file = '/home/kfxw/Development/data/VOC/VOC2012-May-11-data/VOCdevkit/VOC2012/ImageSets/Segmentation/deeplab_trainval_aug.txt'
self.val_list_file = '/home/kfxw/Development/data/VOC/VOC2012-May-11-data/VOCdevkit/VOC2012/ImageSets/Segmentation/deeplab_val.txt'
self.test_list_file = '/home/kfxw/Development/data/VOC/VOC2012-May-11-data/VOCdevkit/VOC2012/ImageSets/Segmentation/deeplab_val.txt'
arg = opt()
dataset_train = VOCTrainDataset(arg, batch_per_gpu=4)
loader_train = torchdata.DataLoader(dataset_train,
batch_size=4,
shuffle=False,
collate_fn=user_scattered_collate,
num_workers=1,
drop_last=True,
pin_memory=False)
iterator_train = iter(loader_train)
batch_data = next(iterator_train)
while True:
img = batch_data[0]['data'].cpu().numpy()
plt.subplot(1, 3, 1)
plt.imshow(img[1, ::-1, :, :].transpose(1, 2, 0) / 255.0 + 0.5)
plt.subplot(1, 3, 2)
plt.imshow(img[1, :, :, :].transpose(1, 2, 0) / 255.0)
plt.subplot(1, 3, 3)
seg = batch_data[0]['seg_label'].cpu().numpy()
plt.imshow(seg[1, :, :])
plt.show()
def main(args):
# Network Builders
builder = VGGModelBuilder()
net_encoder = builder.build_encoder(arch=args.arch_encoder,
weights=args.weights_encoder)
net_decoder = builder.build_decoder(arch=args.arch_decoder,
fc_dim=1024,
num_class=args.num_class,
weights=args.weights_decoder)
crit = nn.CrossEntropyLoss(ignore_index=255, reduction='sum')
if args.arch_decoder.endswith('deepsup'):
segmentation_module = SegmentationModule(net_encoder, net_decoder,
crit, args.deep_sup_scale)
else:
segmentation_module = SegmentationModule(net_encoder, net_decoder,
crit)
print(segmentation_module)
# Dataset and Loader
dataset_train = VOCTrainDataset(args,
batch_per_gpu=args.batch_size_per_gpu)
loader_train = torchdata.DataLoader(
dataset_train,
batch_size=1, # data_parallel have been modified, not useful
shuffle=False, # do not use this param
collate_fn=user_scattered_collate,
num_workers=1, # MUST be 1 or 0
drop_last=True,
pin_memory=True)
print('[{}] 1 training epoch = {} iters'.format(
args.epoch_iters,
datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")))
# create loader iterator
iterator_train = iter(loader_train)
# load nets into gpu
if args.num_gpus > 1:
segmentation_module = UserScatteredDataParallel(segmentation_module,
device_ids=args.gpu_id)
# For sync bn
patch_replication_callback(segmentation_module)
segmentation_module.cuda(device=args.gpu_id[0])
# Set up optimizers
nets = (net_encoder, net_decoder, crit)
optimizers = create_optimizers(nets, args)
# Main loop
history = {
'train': {
'epoch': [],
'loss': [],
'train_acc': [],
'test_ious': [],
'test_mean_iou': []
}
}
for epoch in range(args.start_epoch, args.num_epoch + 1):
# test/validate
dataset_val = VOCValDataset(
args
) # create val dataset loader for every eval, in order to use drop_last=false
loader_val = torchdata.DataLoader(
dataset_val,
# data_parallel have been modified, MUST use val batch size
# and collate_fn MUST be user_scattered_collate
batch_size=args.val_batch_size,
shuffle=False,
collate_fn=user_scattered_collate,
num_workers=1, # MUST be 1 or 0
drop_last=False)
iterator_val = iter(loader_val)
if epoch % args.test_epoch_interval == 0 or epoch == args.num_epoch: # epoch != 1 and
(cls_ious, cls_mean_iou) = evaluate(segmentation_module,
iterator_val, args)
history['train']['test_ious'].append(cls_ious)
history['train']['test_mean_iou'].append(cls_mean_iou)
else:
history['train']['test_ious'].append(-1) # empty data
history['train']['test_mean_iou'].append(-1)
# train
train(segmentation_module, iterator_train, optimizers, history, epoch,
args)
# checkpointing
checkpoint(nets, history, args, epoch)
print('[{}] Training Done!'.format(
datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")))
def main():
builder = ModelBuilder()
net_encoder = builder.build_encoder_textmountain
net_decoder = builder.build_decoder_textmountain
encoder_path = 'model_path'
decoder_path = 'model_path'
net_encoder = net_encoder(encoder_path)
net_decoder = net_decoder(decoder_path)
segmentation_module = Module_textmountain(net_encoder,
net_decoder,
score_thres=0.7)
segmentation_module.cuda()
segmentation_module.eval()
dataset_train = ValDataset(size=test_size)
loader_train = torchdata.DataLoader(
dataset_train,
batch_size=1, # we have modified data_parallel
shuffle=False, # we do not use this param
collate_fn=user_scattered_collate,
num_workers=int(1),
drop_last=True,
pin_memory=True)
iterator_train = iter(loader_train)
for i_data in range(len(iterator_train)):
batch_data = next(iterator_train)
with torch.no_grad():
feed_dict = {}
feed_dict['images'] = batch_data[0]['im'].cuda()
pred_polygon, score_box = segmentation_module.forward(
feed_dict,
is_Train=False,
img_name=batch_data[0]['im_name'],
save_dir=save_dir)
if len(pred_polygon) == 0:
continue
im = batch_data[0]['im_name']
scale_w = batch_data[0]['scale_w']
scale_h = batch_data[0]['scale_h']
pred_polygon[:, :, 0] /= scale_w
pred_polygon[:, :, 1] /= scale_h
if pred_polygon is not None:
res_file = os.path.join(
save_dir, 'res_{}.txt'.format(
os.path.basename(im)[3:].split('.')[0]))
with open(res_file, 'w') as f:
for box, score_one in zip(pred_polygon, score_box):
box = box.astype(np.int32)
if np.linalg.norm(box[0] -
box[1]) < 5 or np.linalg.norm(
box[3] - box[0]) < 5:
continue
f.write('{},{},{},{},{},{},{},{},{}\r\n'.format(
box[0, 0], box[0, 1], box[1, 0], box[1, 1], box[2,
0],
box[2, 1], box[3, 0], box[3, 1], score_one))
def main(args):
torch.cuda.set_device(args.gpu)
# Network Builders
builder = ModelBuilder()
net_encoder = None
net_decoder = None
unet = None
if args.unet == False:
net_encoder = builder.build_encoder(arch=args.arch_encoder,
fc_dim=args.fc_dim,
weights=args.weights_encoder)
net_decoder = builder.build_decoder(arch=args.arch_decoder,
fc_dim=args.fc_dim,
num_class=args.num_class,
weights=args.weights_decoder,
use_softmax=True)
else:
unet = builder.build_unet(num_class=args.num_class,
arch=args.arch_unet,
weights=args.weights_unet,
use_softmax=True)
crit = nn.NLLLoss()
if args.unet == False:
segmentation_module = SegmentationModule(net_encoder, net_decoder,
crit)
else:
segmentation_module = SegmentationModule(net_encoder,
net_decoder,
crit,
is_unet=args.unet,
unet=unet)
'''
# Dataset and Loader
dataset_val = dl.loadVal()
loader_val = torchdata.DataLoader(
dataset_val,
batch_size=5,
shuffle=False,
num_workers=1,
drop_last=True)
'''
# Dataset and Loader
dataset_val = ValDataset(args.list_val, args, max_sample=args.num_val)
loader_val = torchdata.DataLoader(dataset_val,
batch_size=args.batch_size,
shuffle=False,
collate_fn=user_scattered_collate,
num_workers=5,
drop_last=True)
segmentation_module.cuda()
# Main loop
evaluate(segmentation_module, loader_val, args)
print('Evaluation Done!')
def main(args):
torch.cuda.set_device(args.gpu)
# Network Builders
builder = ModelBuilder()
net_encoder = builder.build_encoder(arch=args.arch_encoder,
fc_dim=args.fc_dim,
weights=args.weights_encoder)
net_decoder = builder.build_decoder(arch=args.arch_decoder,
fc_dim=args.fc_dim,
num_class=args.num_class,
weights=args.weights_decoder,
use_softmax=True)
crit = nn.NLLLoss(ignore_index=-1)
segmentation_module = SegmentationModule(net_encoder, net_decoder, crit)
# Dataset and Loader
# list_test = [{'fpath_img': args.test_img}]
# test_chk = []
# testing = os.listdir("/home/teai/externalhd2/BDD100K/segmentation_v2/test/")
# for i in testing:
# if(i.endswith(".jpg")):
# test_chk.append("/home/teai/externalhd2/BDD100K/segmentation_v2/test/"+i)
video_path = "./test_video_input/test_1.mp4"
vidcap = cv2.VideoCapture(video_path)
video_fps = math.ceil(vidcap.get(cv2.CAP_PROP_FPS))
length = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT))
frame_array = []
for i in tqdm(range(length - 1)):
ret, frame = vidcap.read()
cv2.imwrite("./test_video_input/frame.png", frame)
test_chk = ['./test_video_input/frame.png']
# print(type(args.test_imgs))
list_test = [{'fpath_img': x} for x in test_chk]
#list_test=[{'fpath_img': 'frame_143.png'},{'fpath_img': 'frame_100.png'},{'fpath_img': 'frame_1.png'}]
#print("list_test",list_test)
dataset_test = TestDataset(list_test, args, max_sample=args.num_val)
loader_test = torchdata.DataLoader(dataset_test,
batch_size=args.batch_size,
shuffle=False,
collate_fn=user_scattered_collate,
num_workers=5,
drop_last=True)
segmentation_module.cuda()
# Main loop
# start=time.time()
test(segmentation_module, loader_test, args)
# end=time.time()
# print("Time taken",(end-start))
#print('Inference done!')
img = cv2.imread("./test_video_output/frame.png")
height, width, layers = img.shape
size = (width, height)
frame_array.append(img)
out = cv2.VideoWriter("./test_video_output/test_1_sgd_100.mp4",
cv2.VideoWriter_fourcc(*'DIVX'), video_fps, size)
for i in range(len(frame_array)):
# writing to a image array
out.write(frame_array[i])
out.release()
# 2. network setup
builder = VGGModelBuilder()
net_encoder = builder.build_encoder(arch=args.arch_encoder,
weights=args.weights_encoder)
net_decoder = builder.build_decoder(arch=args.arch_decoder,
num_class=args.num_class,
fc_dim=1024,
weights=args.weights_decoder)
crit = nn.CrossEntropyLoss(ignore_index=255)
segmentation_module = SegmentationModule(net_encoder, net_decoder, crit)
# 3. dataset
dataset_val = VOCValDataset(args)
loader_val = torchdata.DataLoader(dataset_val,
batch_size=args.val_batch_size,
shuffle=False,
collate_fn=user_scattered_collate,
num_workers=1,
drop_last=False)
iterator_val = iter(loader_val)
nets = (net_encoder, net_decoder, crit)
segmentation_module.cuda(device=args.gpu_id[0])
# 5. validate
def evaluate(segmentation_module, loader, args):
acc_meter = AverageMeter()
intersection_meter = AverageMeter()
union_meter = AverageMeter()
cls_ious_meter = AverageMeter()
cls_mean_iou_meter = AverageMeter()
def main(args):
# Network Builders
builder = ModelBuilder()
net_encoder = builder.build_encoder(arch=args.arch_encoder,
fc_dim=args.fc_dim,
weights=args.weights_encoder)
net_decoder = builder.build_decoder(arch=args.arch_decoder,
fc_dim=args.fc_dim,
num_class=args.num_class,
weights=args.weights_decoder)
crit = nn.NLLLoss(ignore_index=255)
if args.arch_decoder.endswith('deepsup'):
segmentation_module = SegmentationModule(net_encoder, net_decoder,
crit, args.deep_sup_scale)
else:
segmentation_module = SegmentationModule(net_encoder, net_decoder,
crit)
# Dataset and Loader
# dataset_train = TrainDataset(
# args.list_train, args, batch_per_gpu=args.batch_size_per_gpu)
# dataset_train = voc.TrainDataset_VOC(dataset_root=config.img_root_folder, mode='train', transform=input_transform)
dataset_train = VOC_TrainDataset(opt=args,
batch_per_gpu=args.batch_size_per_gpu)
loader_train = torchdata.DataLoader(
dataset_train,
batch_size=args.num_gpus, # we have modified data_parallel
shuffle=False, # we do not use this param
collate_fn=user_scattered_collate,
num_workers=int(args.workers),
drop_last=True,
pin_memory=True)
print('1 Epoch = {} iters'.format(args.epoch_iters))
# create loader iterator
iterator_train = iter(loader_train)
# load nets into gpu
if args.num_gpus > 1:
segmentation_module = UserScatteredDataParallel(segmentation_module,
device_ids=range(
args.num_gpus))
# segmentation_module = UserScatteredDataParallel(
# segmentation_module,
# device_ids=[0,3,4,5])
# For sync bn
patch_replication_callback(segmentation_module)
segmentation_module.cuda()
# Set up optimizers
nets = (net_encoder, net_decoder, crit)
optimizers = create_optimizers(nets, args)
if args.resume:
file_path_history = '{}/history_{}'.format(args.ckpt, 'epoch_last.pth')
file_path_encoder = '{}/encoder_{}'.format(args.ckpt, 'epoch_last.pth')
file_path_decoder = '{}/decoder_{}'.format(args.ckpt, 'epoch_last.pth')
file_path_optimizers = '{}/optimizers_{}'.format(
args.ckpt, 'epoch_last.pth')
if os.path.isfile(file_path_history):
print("=> loading checkpoint '{}'".format(file_path_history))
checkpoint_history = torch.load(file_path_history)
checkpoint_encoder = torch.load(file_path_encoder)
checkpoint_decoder = torch.load(file_path_decoder)
checkpoint_optimizers = torch.load(file_path_optimizers)
args.start_epoch = int(checkpoint_history['train']['epoch'][0]) + 1
nets[0].load_state_dict(checkpoint_encoder)
nets[1].load_state_dict(checkpoint_decoder)
optimizers[0].load_state_dict(
checkpoint_optimizers['encoder_optimizer'])
optimizers[1].load_state_dict(
checkpoint_optimizers['decoder_optimizer'])
print("=> start train epoch {}".format(args.start_epoch))
else:
print('resume not find epoch-last checkpoint')
# Main loop
history = {'train': {'epoch': [], 'loss': [], 'acc': []}}
for epoch in range(args.start_epoch, args.num_epoch + 1):
train(segmentation_module, iterator_train, optimizers, history, epoch,
args)
# checkpointing
checkpoint(nets, optimizers, history, args, epoch)
print('Training Done!')