type=str,
default='../../../Output/Spine_Location/latest.pkl')
parser.add_argument('--model_type',
type=str,
default='Unet_base',
help='the type of Unet(Unet_base or Unet_large)')
parser.add_argument('--seed', type=int, default=68, help='set the seed')
args = parser.parse_args()
torch.manual_seed(seed=args.seed)
random.seed(args.seed)
np.random.seed(args.seed)
# Start training
trainer = NetworkTrainer()
trainer.setting.project_name = 'Spine_Location'
trainer.setting.output_dir = '../../../Output/Spine_Location'
list_GPU_ids = args.list_GPU_ids
csv_path = '../../Catalogue' + '/' + str(args.catalogue) + '.csv'
catalogue = csv_to_catalogue(csv_path)
# setting.network is an object
if args.model_type == 'Unet_base':
trainer.setting.network = Model(in_ch=1,
out_ch=1,
list_ch=[-1, 16, 32, 64, 128, 256])
print('Loading Unet_base !')
else:
trainer.setting.network = Model(in_ch=1,
out_ch=1,
from model import Model
from online_evaluation import online_evaluation
from loss import Loss
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--batch_size', type=int, default=32,
help='batch size for training (default: 32)')
parser.add_argument('--list_GPU_ids', nargs='+', type=int, default=[0],
help='list_GPU_ids for training (default: [0])')
parser.add_argument('--max_iter', type=int, default=100000,
help='training iterations(default: 100000)')
args = parser.parse_args()
# Start training
trainer = NetworkTrainer()
trainer.setting.project_name = 'DCNN'
trainer.setting.output_dir = '../../Output/DCNN'
list_GPU_ids = args.list_GPU_ids
trainer.setting.network = Model(in_ch=4, out_ch=1,
list_ch=[-1, 32, 64, 128, 256])
trainer.setting.max_iter = args.max_iter
trainer.setting.train_loader, trainer.setting.val_loader = get_loader(
train_bs=args.batch_size,
val_bs=1,
train_num_samples_per_epoch=args.batch_size * 5000, # 5000 iterations per epoch
val_num_samples_per_epoch=1,
num_works=8