sleep-staging/sc_config.py at master · weberpals/sleep-staging · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
import os
import numpy as np

class Config(object):

    @staticmethod
    def get(model_name):
        if model_name[0:3] == 'oct':
            return OCConfig(model_name)
        elif model_name[0:2] == 'ac':
            return ACConfig(model_name)
        else:
            raise Exception

    def __init__(self, scope, num_features, num_hidden, segsize, lstm, num_classes, batch_size, max_train_len, atonce, restart=True, model_name='small_lstm', is_train=False):

        # Model folder
        root_python = os.path.dirname(os.path.realpath(__file__))
        root_base = str.join('/', root_python.split('/')[:-1])

        self.model_dir = os.path.join(root_python, 'model', scope, model_name)
        self.model_dir_test = os.path.join(root_python, 'model', scope, model_name)
        if not os.path.isdir(self.model_dir):
            os.mkdir(self.model_dir)

        # Data
        data_dir = '/scratch/users/jenss/'

	if model_name[0]=='o':
	        self.train_data = os.path.join(data_dir, 'octave_training_data')
		self.test_data = os.path.join(data_dir, 'octave_test_data2')

        else:
		self.train_data = os.path.join(data_dir, 'ac_training_data')
		self.test_data = os.path.join(data_dir, 'ac_test_data')

        self.train_dir = os.path.join(self.model_dir, 'train')
        if not os.path.isdir(self.train_dir):
            os.mkdir(self.train_dir)

        # Configuration
        self.model_name = model_name
        self.scope = scope
        self.load_list = self.model_dir+'_load_list.csv'
        self.load_prob_file = self.model_dir+'_load_list.h5'
        self.validation_results = self.model_dir+'_training_data.csv'
        self.is_training = is_train
        self.num_features = num_features
        self.num_classes = num_classes
        self.batch_size = batch_size
        self.restart = restart
        self.lstm = lstm
        self.num_hidden = num_hidden
        self.keep_prob = 0.5
        self.segsize = segsize
        self.eval_nseg_atonce = atonce
        self.max_train_len = max_train_len
	self.save_freq = 1

        # Training
        self.max_steps = 500000

    def checkpoint_file(self, ckpt=0):
        if ckpt == 0:
            return os.path.join(self.model_dir, 'model.ckpt')
        else:
            return os.path.join(self.model_dir, 'model.ckpt-%.0f' % ckpt)


class OCConfig(Config):

    def __init__(self,restart=True, model_name='oct_sh_ss_lstm', is_training=False):
        print('model: '+model_name)
        scope = 'oct'
        num_features = 25
        num_classes = 5

        max_train_len = 360000


        if is_training:
            batch_size = 15
        else:
            batch_size = 1

        if model_name[4:6]=='lh':
            num_hidden = 256
        else:
            num_hidden = 128
        if model_name[7:9]=='ls':
            segsize = 1500
	    atonce = 1500
        else:
            segsize = 500
	    atonce = 4500
        if model_name[10:12]=='ff':
            lstm = False
        else:
            lstm = True


        super(OCConfig, self).__init__(scope, num_features, num_hidden, segsize, lstm, num_classes, batch_size, max_train_len, atonce, restart, model_name, is_training)

class ACConfig(Config):

    def __init__(self,restart=True, model_name='ac_sh_ss_lstm', is_training=False):
        print('model: '+model_name)
        scope = 'ac'
        num_features = 1640
        num_classes = 5
	if is_training:
        	batch_size = 5
	else:
		batch_size = 1

        max_train_len = 14400

        if model_name[3:5]=='lh':
            num_hidden = 256
	elif model_name[3:5]=='rh':
	    	np.random.seed(int(model_name[-2:]))
	    	num_hidden = 256+np.round(np.random.rand(1)*128)
		num_hidden = num_hidden[0].astype(int)
        else:
            num_hidden = 128
        if model_name[6:8]=='ls':
            segsize = 60
	    atonce = 700
        else:
            segsize = 20
	    atonce = 3000
        if model_name[9:11]=='ff':
            lstm = False
        else:
            lstm = True


        super(ACConfig, self).__init__(scope, num_features, num_hidden, segsize, lstm, num_classes, batch_size, max_train_len, atonce, restart, model_name, is_training)