iframe-proxy | Sunbelt Computer Software

This repository was archived by the owner on Jun 11, 2026. It is now read-only.

Repository navigation

History

# Copyright (c) Microsoft Corporation.

# Licensed under the MIT License.

'''Training script for DIF-Net.

'''

import sys

import os

sys.path.append( os.path.dirname( os.path.dirname( os.path.abspath(__file__) ) ) )

import yaml

import io

import numpy as np

import dataset, utils, training_loop, loss, modules, meta_modules

import torch

from torch.utils.data import DataLoader

import configargparse

from torch import nn

from dif_net import DeformedImplicitField

os.environ['CUDA_VISIBLE_DEVICES'] = '0,1,2,3,4,5,6,7'

p = configargparse.ArgumentParser()

p.add_argument('--config', type=str,default='', help='training configuration.')

p.add_argument('--train_split', type=str,default='', help='training subject names.')

p.add_argument('--logging_root', type=str, default='./logs', help='root for logging')

p.add_argument('--experiment_name', type=str, default='default',

help='Name of subdirectory in logging_root where summaries and checkpoints will be saved.')

# General training options

p.add_argument('--batch_size', type=int, default=256, help='training batch size.')

p.add_argument('--lr', type=float, default=1e-4, help='learning rate. default=1e-4')

p.add_argument('--epochs', type=int, default=500, help='Number of epochs to train for.')

p.add_argument('--epochs_til_checkpoint', type=int, default=5,

help='Time interval in seconds until checkpoint is saved.')

p.add_argument('--steps_til_summary', type=int, default=100,

help='Time interval in seconds until tensorboard summary is saved.')

p.add_argument('--model_type', type=str, default='sine',

help='Options are "sine" (all sine activations) and "mixed" (first layer sine, other layers tanh)')

p.add_argument('--point_cloud_path', type=str, default='',

help='training data path.')

p.add_argument('--latent_dim', type=int,default=128, help='latent code dimension.')

p.add_argument('--hidden_num', type=int,default=128, help='hidden layer dimension of deform-net.')

p.add_argument('--loss_grad_deform', type=float,default=5, help='loss weight for deformation smoothness prior.')

p.add_argument('--loss_grad_temp', type=float,default=1e2, help='loss weight for normal consistency prior.')

p.add_argument('--loss_correct', type=float,default=1e2, help='loss weight for minimal correction prior.')

p.add_argument('--num_instances', type=int,default=5, help='numbers of instance in the training set.')

p.add_argument('--expand', type=float,default=-1, help='expansion of shape surface.')

p.add_argument('--max_points', type=int,default=200000, help='number of surface points for each epoch.')

p.add_argument('--on_surface_points', type=int,default=4000, help='number of surface points for each iteration.')

# load configs if exist

opt = p.parse_args()

if opt.config == '':

meta_params = vars(opt)

else:

with open(opt.config,'r') as stream:

meta_params = yaml.safe_load(stream)

# define dataloader

with open(meta_params['train_split'],'r') as file:

all_names = file.read().split('\n')

data_path = [os.path.join(meta_params['point_cloud_path'],f + '.mat') for f in all_names]

sdf_dataset = dataset.PointCloudMulti(root_dir=data_path, max_num_instances=meta_params['num_instances'],**meta_params)

dataloader = DataLoader(sdf_dataset, shuffle=True,collate_fn=sdf_dataset.collate_fn,batch_size=meta_params['batch_size'], pin_memory=True, num_workers=16, drop_last = True)

print('Total subjects: ',sdf_dataset.num_instances)

meta_params['num_instances'] = sdf_dataset.num_instances

# define DIF-Net

model = DeformedImplicitField(**meta_params)

model = nn.DataParallel(model)

model.cuda()

# create save path

root_path = os.path.join(meta_params['logging_root'], meta_params['experiment_name'])

utils.cond_mkdir(root_path)

with io.open(os.path.join(root_path,'model.yml'),'w',encoding='utf8') as outfile:

yaml.dump(meta_params, outfile, default_flow_style=False, allow_unicode=True)

# main training loop

training_loop.train(model=model, train_dataloader=dataloader, model_dir=root_path, **meta_params)

Sunbelt Computer Software

PL/B Language Development and Support

Navigation Menu

Repository navigation

train.py

train.py

Sunbelt Computer Software

PL/B Language Development and Support

FilesExpand file tree

train.py

Latest commit

History

train.py

File metadata and controls