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):
# 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!')