-
Notifications
You must be signed in to change notification settings - Fork 1
/
train.py
138 lines (109 loc) · 4.18 KB
/
train.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
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
#!/usr/bin/env python
# coding: utf-8
import time
import numpy as np
import torch
import h5py
import functools
import math
from torch import nn, optim
from torch.autograd import Variable
import torch.nn.functional as F
from torch.optim import lr_scheduler
from torch.utils.data import Dataset
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt
import os
from datetime import datetime
import parameters
from scipy import io
import sys
sys.path.append("..")
print('current time:',datetime.now())
## GPU
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print(device)
## import net
from source_transformation import wavenet2d as myNet
net = myNet()
def SNR(noisy,gt):
res = noisy - gt
msegt = np.mean(gt * gt)
mseres = np.mean(res * res)
SNR = 10 * math.log((msegt/mseres),10)
return SNR
## load data
sample_size_train = parameters.sample_size_train
sample_size_test = parameters.sample_size_test
x = np.zeros([sample_size_train, 1,parameters.timespan_input, parameters.trace])
y = np.zeros([sample_size_train, 1,parameters.timespan, parameters.trace])
X = np.zeros([sample_size_test,1,parameters.timespan_input,parameters.trace])
Y = np.zeros([sample_size_test,1,parameters.timespan,parameters.trace])
f = h5py.File(parameters.data_path, 'r')
x[:,:,:parameters.timespan,:] = f['X'][0:sample_size_train,:,:]
y[:,:,:,:] = f['Y'][0:sample_size_train,:,:]
f.close()
f = h5py.File(parameters.test_data_path, 'r')
X[:,:,:parameters.timespan,:] = f['X'][0:sample_size_test,:,:]
Y[:,:,:,:] = f['Y'][0:sample_size_test,:,:]
f.close()
## parameter
class MyDataset(Dataset):
def __init__(self, a, b):
self.data_1 = a
self.data_2 = b
def __len__(self):
return len(x)
def __getitem__(self, idx):
in_put = self.data_1[idx]
out_put = self.data_2[idx]
return in_put, out_put
batchsize = parameters.batchsize
dataset = MyDataset(x,y)
train_iter = DataLoader(dataset, batch_size=batchsize, shuffle=True, num_workers=10, drop_last=False, pin_memory=True)
lr = parameters.learning_rate
num_epochs = parameters.num_epochs
optimizer = torch.optim.Adam(net.parameters(), lr = lr, betas=(0.9, 0.999), eps=1e-8)
scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=num_epochs, eta_min=parameters.end_lr)
loss_res = np.zeros(num_epochs)
valida_res = np.zeros(num_epochs)
if parameters.checkpoint_epoch>0:
net.load_state_dict(torch.load(parameters.result_path+str(parameters.checkpoint_epoch)+'.pkl'))
net = net.to(device)
print("training on ", device)
loss = torch.nn.L1Loss(reduction='sum') #torch.nn.MSELoss(reduction='sum')
## load test data in GPU
Xt = Variable(torch.from_numpy(X))
Xt = Xt.to(device)
Xt = Xt.type(torch.cuda.FloatTensor)
for epoch in range(num_epochs):
train_l_sum = 0.0
start = time.time()
batch_count = 0
for xtrain, ytrain in train_iter:
xtrain = xtrain.to(device).type(torch.cuda.FloatTensor)
ytrain = ytrain.to(device).type(torch.cuda.FloatTensor)
y_hat = net(xtrain)
l = loss(torch.squeeze(y_hat), torch.squeeze(ytrain))
optimizer.zero_grad()
l.backward()
optimizer.step()
train_l_sum += l.cpu().item()
batch_count += 1
scheduler.step()
with torch.no_grad():
Y_hat = net(Xt)
Y_hat = Y_hat.data.cpu().numpy()
Y_hat = Y_hat.reshape(sample_size_test,1,parameters.timespan,parameters.trace)
snr = np.mean(SNR(Y_hat,Y))
print('epoch %d, loss %.6f, validation %.6f, time %.1f sec'
% (epoch +parameters.checkpoint_epoch+ 1, train_l_sum/batch_count , snr, time.time() - start))
loss_res[epoch] = train_l_sum/batch_count
valida_res[epoch] = snr
if ((epoch+1) % 5) == 0:
torch.save(net.state_dict(), parameters.result_path+str(epoch+parameters.checkpoint_epoch+1)+'.pkl')
io.savemat(parameters.result_path+'training_epoch.mat',{'loss_res':loss_res,'valida_res':valida_res})
np.savetxt('loss.csv', loss_res, delimiter = '')
np.savetxt('validation.csv', valida_res, delimiter = '')
print('smallest error on trainging set:',np.argmin(loss_res)+1,'smallest error on test set:',np.argmax(valida_res)+1)
print('current time:',datetime.now())