def main(cfg, gpus):
# Network Builders
net_encoder = ModelBuilder.build_encoder(
arch=cfg.MODEL.arch_encoder.lower(),
fc_dim=cfg.MODEL.fc_dim,
weights=cfg.MODEL.weights_encoder)
net_decoder = ModelBuilder.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)
if cfg.MODEL.arch_decoder.endswith('regression'):
crit = nn.MSELoss(
reduction="sum"
) # Sum for multi-output learning, need to sum across all labels
else:
crit = nn.NLLLoss(ignore_index=-1) # negative log likelihood loss
if cfg.MODEL.arch_decoder.endswith('deepsup'):
segmentation_module = SegmentationModule(net_encoder, net_decoder,
crit, cfg.DATASET.classes,
cfg.TRAIN.deep_sup_scale)
else:
segmentation_module = SegmentationModule(net_encoder, net_decoder,
crit, cfg.DATASET.classes)
print("net_encoder")
print(type(net_encoder))
print(net_encoder)
print("net_decoder")
print(type(net_decoder))
print(net_decoder)
# Dataset and Loader
if cfg.MODEL.arch_decoder.endswith('regression'):
print("performing regression")
dataset_train = TrainDatasetRegression(
cfg.DATASET.root_dataset,
cfg.DATASET.list_train,
cfg.DATASET.classes,
cfg.DATASET,
batch_per_gpu=cfg.TRAIN.batch_size_per_gpu)
dataset_val = TrainDatasetRegression(
cfg.DATASET.root_dataset,
cfg.DATASET.list_val,
cfg.DATASET.classes,
cfg.DATASET,
batch_per_gpu=cfg.TRAIN.batch_size_per_gpu)
else:
dataset_train = TrainDataset(
cfg.DATASET.root_dataset,
cfg.DATASET.list_train,
cfg.DATASET,
batch_per_gpu=cfg.TRAIN.batch_size_per_gpu)
loader_train = torch.utils.data.DataLoader(
dataset_train,
batch_size=len(gpus), # we have modified data_parallel
shuffle=False, # we do not use this param
collate_fn=user_scattered_collate,
num_workers=cfg.TRAIN.workers,
drop_last=True,
pin_memory=True)
loader_val = torch.utils.data.DataLoader(
dataset_val,
batch_size=len(gpus), # we have modified data_parallel
shuffle=False, # we do not use this param
collate_fn=user_scattered_collate,
num_workers=cfg.TRAIN.workers,
drop_last=True,
pin_memory=True)
print('1 Epoch = {} iters'.format(cfg.TRAIN.epoch_iters))
# create loader iterator
iterator_train = iter(loader_train)
iterator_val = iter(loader_val)
# load nets into gpu
if len(gpus) > 1:
segmentation_module = UserScatteredDataParallel(segmentation_module,
device_ids=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, cfg)
# Main loop
history = {
'train': {
'epoch': [],
'loss': []
},
'val': {
'epoch': [],
'loss': []
}
}
for epoch in range(cfg.TRAIN.start_epoch, cfg.TRAIN.num_epoch):
train(segmentation_module, iterator_train, optimizers, history,
epoch + 1, cfg)
val(segmentation_module, iterator_val, optimizers, history, epoch + 1,
cfg)
# checkpointing every 5th epoch
if (epoch % 5 == 0):
checkpoint(nets, history, cfg, epoch + 1)
print('Training Done!')
def main(cfg, gpus):
# Network Builders
net_encoder = ModelBuilder.build_encoder(
arch=cfg.MODEL.arch_encoder.lower(),
fc_dim=cfg.MODEL.fc_dim,
weights=cfg.MODEL.weights_encoder)
net_decoder = ModelBuilder.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)
crit = nn.NLLLoss(ignore_index=-1)
if cfg.MODEL.arch_decoder.endswith('deepsup'):
segmentation_module = SegmentationModule(net_encoder, net_decoder,
crit,
cfg.TRAIN.deep_sup_scale)
else:
segmentation_module = SegmentationModule(net_encoder, net_decoder,
crit)
# Dataset and Loader
dataset_train = TrainDataset(cfg.DATASET.root_dataset,
cfg.DATASET.list_train,
cfg.DATASET,
batch_per_gpu=cfg.TRAIN.batch_size_per_gpu)
loader_train = torch.utils.data.DataLoader(
dataset_train,
batch_size=len(gpus), # we have modified data_parallel
shuffle=False, # we do not use this param
collate_fn=user_scattered_collate,
num_workers=cfg.TRAIN.workers,
drop_last=True,
pin_memory=True)
print('1 Epoch = {} iters'.format(cfg.TRAIN.epoch_iters))
# create loader iterator
iterator_train = iter(loader_train)
# load nets into gpu
print(gpus)
# if len(gpus) > 1:
if True:
segmentation_module = UserScatteredDataParallel(segmentation_module,
device_ids=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, cfg)
# Main loop
history = {'train': {'epoch': [], 'loss': [], 'acc': []}}
for epoch in range(cfg.TRAIN.start_epoch, cfg.TRAIN.num_epoch):
train(segmentation_module, iterator_train, optimizers, history,
epoch + 1, cfg)
# checkpointing
checkpoint(nets, history, cfg, epoch + 1)
print('Training Done!')
collate_fn=user_scattered_collate,
num_workers=cfg.TRAIN.workers,
drop_last=True,
pin_memory=True)
print('1 Epoch = {} iters'.format(cfg.TRAIN.epoch_iters))
# create loader iterator
iterator_train = iter(loader_train)
# load nets into gpu
if len(gpus) > 1:
segmentation_module = UserScatteredDataParallel(
segmentation_module,
device_ids=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, cfg)
# Main loop
history = {'train': {'epoch': [], 'loss': [], 'acc': []}}
for epoch in range(cfg.TRAIN.start_epoch, cfg.TRAIN.num_epoch):
train(segmentation_module, iterator_train, optimizers, history, epoch+1, cfg)
# checkpointing
checkpoint(nets, history, cfg, epoch+1)