# model, optimizer
model = NeuralRecon(cfg)
if cfg.DISTRIBUTED:
model.cuda()
model = DistributedDataParallel(
model,
device_ids=[cfg.LOCAL_RANK],
output_device=cfg.LOCAL_RANK,
# this should be removed if we update BatchNorm stats
broadcast_buffers=False,
find_unused_parameters=True)
else:
model = torch.nn.DataParallel(model, device_ids=[0])
model.cuda()
optimizer = torch.optim.Adam(model.parameters(),
lr=cfg.TRAIN.LR,
betas=(0.9, 0.999),
weight_decay=cfg.TRAIN.WD)
# main function
def train():
# load parameters
start_epoch = 0
if cfg.RESUME:
saved_models = [
fn for fn in os.listdir(cfg.LOGDIR) if fn.endswith(".ckpt")
]
saved_models = sorted(
saved_models, key=lambda x: int(x.split('_')[-1].split('.')[0]))
drop_last=False)
# model, optimizer
model = NeuralRecon(cfg)
if cfg.DISTRIBUTED:
model.cuda()
model = DistributedDataParallel(
model, device_ids=[cfg.LOCAL_RANK], output_device=cfg.LOCAL_RANK,
# this should be removed if we update BatchNorm stats
broadcast_buffers=False,
find_unused_parameters=True
)
else:
model = torch.nn.DataParallel(model, device_ids=[0])
model.cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=cfg.TRAIN.LR, betas=(0.9, 0.999), weight_decay=cfg.TRAIN.WD)
# main function
def train():
# load parameters
start_epoch = 0
if cfg.RESUME:
saved_models = [fn for fn in os.listdir(cfg.LOGDIR) if fn.endswith(".ckpt")]
saved_models = sorted(saved_models, key=lambda x: int(x.split('_')[-1].split('.')[0]))
if len(saved_models) != 0:
# use the latest checkpoint file
loadckpt = os.path.join(cfg.LOGDIR, saved_models[-1])
logger.info("resuming " + str(loadckpt))
map_location = {'cuda:%d' % 0: 'cuda:%d' % cfg.LOCAL_RANK}
state_dict = torch.load(loadckpt, map_location=map_location)