|
- from models.hysnet.solver import HySNet_Solver
-
- import torch
- import argparse
- import os
- import numpy as np
- from torchsummary import summary
- import torch.multiprocessing as mp
- import pandas as pd
-
-
-
- def test_from_csv(dir_path):
- csv_path = os.path.join(dir_path, 'score_record.csv')
- df = pd.read_csv(csv_path)
- for key in df.keys():
- print('mean ', key, ': ', np.mean(df[key]))
- print('val ', key, ': ', np.var(df[key]))
-
- def parse_args():
- parser = argparse.ArgumentParser(description='train')
- # config
- parser.add_argument('--ckp-path', default=None, type=str,
- help='path of checkpoint, .pth or None')
- parser.add_argument('--output-dir', default='output-0', type=str,
- help='path of output dir')
- parser.add_argument('--dataset', default='mriqc', choices=['mriqc', 'mriqc2d'],
- help='Dataset, mriqc or mriqc2d', type=str)
- parser.add_argument('--train-path', default='/home/qkh/data/MRIQC/T1w/', type=str,
- help='path of train dir')
- parser.add_argument('--val-path', default='/home/qkh//data/MRIQC/T1w/', type=str,
- help='path of validation dir')
- parser.add_argument('--test-path', default='/home/qkh/data/MRIQC/T1w/', type=str,
- help='path of test dir')
-
- parser.add_argument('-g', '--gpu', default=0, type=int)
- # parser.add_argument('-g', '--gpu', default='0', type=str)
- parser.add_argument('--mode', default='train', type=str, choices=['train', 'test', 'predict'])
- # Model Setting
- parser.add_argument('--model', default='hysnet', type=str, choices=['hysnet'])
- parser.add_argument('--pooling', default='gmaxp', type=str, choices=['gap', 'gmaxp', 'gminp', 'gmaxminp'])
- parser.add_argument('--arch', default='resnet50', type=str)
- # Train Setting
- parser.add_argument('--batch-size', default=8, type=int,
- help='batch size')
- parser.add_argument('--patch-size', default=96, type=int,
- help='batch size')
- parser.add_argument('--patch-num', default=10, type=int,
- help='batch size')
- parser.add_argument('--start-epoch', default=0, type=int)
- parser.add_argument('--num-epochs', default=50, type=int,
- help='number of epochs to train, default: 50')
- parser.add_argument('--lr', type=float, default=0.0001, metavar='LR',
- help='learning rate (default: 0.0001)')
-
- args = parser.parse_args()
-
- print('-' * 20)
- for key in args.__dict__:
- print(key, '=', args.__dict__[key])
- print('-' * 20)
-
- return args
-
-
- def main(args=None):
- if args.model == 'hysnet':
- solver = HySNet_Solver(args=args)
-
- if args.mode == 'train':
- solver.train()
- if args.mode == 'test':
- solver.test()
- if args.mode == 'predict':
- solver.predict()
-
-
- if __name__ == '__main__':
- args = parse_args()
- torch.cuda.set_device(args.gpu)
- # os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
- main(args=args)
|