-
Notifications
You must be signed in to change notification settings - Fork 11
/
main.py
85 lines (80 loc) · 2.7 KB
/
main.py
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
import numpy as np
from train import train
from loader import LandslideDataset, DistLandslideDataset
from torch.utils.data import DataLoader
# from dimension_reduction import reduce_dim
from time import ctime
from sacred import Experiment
# from sacred.observers import MongoObserver
ex = Experiment('CNNPatch')
@ex.config
def ex_cfg():
train_param = {
'optim': 'Adam',
'lr': 0.0001,
'n_epochs': 100,
'bs': 4,
'decay': 1e-3,
'patience': 2,
'pos_weight': 1,
'model': 'UNet'
}
data_param = {
'n_workers': 4,
'region': 'Veneto',
'pix_res': 10,
'stride': 500,
'ws': 500,
'pad': 64,
'feature_num': 94,
'oversample': False,
'prune': 64,
'dist_num': 3, #corresponding to 30,100,300
'dist_feature': False
}
loc_param = {
'load_model': '',
'data_path': '/tmp/Veneto_data.h5',
'index_path': '/home/ainaz/Projects/Landslides/image_data/new_partitioning/',
'save': 20
}
# def plot_grid(x, y):
# import matplotlib.pyplot as plt
# fig = plt.figure()
# fig.add_subplot(1,1,1)
# plt.scatter(x, y['Adam'], c='b')
# plt.scatter(x, y['SGD'], c='r')
# plt.show()
@ex.automain
def main(train_param, data_param, loc_param, _log, _run):
data = []
if data_param['dist_feature']:
for flag in ['train', 'validation']:
data.append(
DistLandslideDataset(
loc_param['data_path'],
np.load(loc_param['index_path']+'{}_{}_indices.npy'.format(data_param['region'], flag)),
data_param['region'],
data_param['ws'],
data_param['pad'],
data_param['prune'],
data_param['dist_num']
)
)
else:
for flag in ['train', 'validation']:
data.append(
LandslideDataset(
loc_param['data_path'],
np.load(loc_param['index_path']+'{}_{}_indices.npy'.format(data_param['region'], flag)),
data_param['region'],
data_param['ws'],
data_param['pad'],
data_param['prune']
)
)
loader = [DataLoader(d, batch_size=train_param['bs'], shuffle=True, num_workers=data_param['n_workers']) for d in data]
_log.info('[{}]: created train and validation datasets.'.format(ctime()))
_log.info('[{}]: starting to train ...'.format(ctime()))
train(loader[0], loader[1], train_param, data_param, loc_param, _log, _run)
_log.info('[{}]: training is finished!'.format(ctime()))