-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathRegressor.py
105 lines (90 loc) · 3.7 KB
/
Regressor.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
#!/usr/bin/env python
#SBATCH --partition=p40_4,p100_4
#SBATCH --nodes=1
#SBATCH --gres=gpu:2
#SBATCH --time=02:00:00
#SBATCH --mem=5GB
#SBATCH --cpus-per-task=14
#SBATCH -o "/scratch/dsw310/CCA/output/Reg/outReg.log"
#SBATCH -e "/scratch/dsw310/CCA/output/Reg/errReg.log"
#SBATCH --mail-type=ALL
#SBATCH [email protected]
#Look: LR, input net
#Changed: Loss criterion
#Notes: 19s
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torch.utils.data.dataset import Dataset
import sys, time
sys.path.insert(0, '/home/dsw310/CCA/functions')
from uNet import *
from tools import *
from data_utils import *
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
dire='/scratch/dsw310/CCA/output/Reg/'
num_epochs=1500
batch_size=28; num_workers=14
#eval_frequency=2
IndList2=np.loadtxt('/scratch/dsw310/CCA/data/extras/IndList/2phase_1.75Thres.dat').astype(int)
class SimuData2(Dataset):
def __init__(self,lIndex,hIndex,aug=0,test=0):
self.datafiles = []
self.aug=aug
self.test=test
for i in np.arange(lIndex,hIndex):
self.datafiles+=[IndList2[i]]
def __getitem__(self, index):
return get_mini_batch(self.datafiles[index],self.aug,self.test)
def __len__(self):
return len(self.datafiles)
best_val = 1e10; thres=1.
mask_model = nn.DataParallel(DMUnet(BasicBlock).to(device))
mask_model.load_state_dict(torch.load('/scratch/dsw310/CCA/Saved/BestModel/Unet/0325.pt'))
mask_model.eval()
for param in mask_model.parameters(): param.requires_grad = False
reg_model = nn.DataParallel(RegUnet(BasicBlock).to(device))
reg_model.load_state_dict(torch.load('/scratch/dsw310/CCA/Saved/BestModel/Reg/0422.pt'))
criterion = reg_loss(weight_ratio=0.1,thres=thres)
optimizer = torch.optim.Adam(reg_model.parameters(),lr=1e-9, betas=(0.9, 0.999), eps=1e-08,weight_decay=1e-3)
start_time = time.time()
TrainSet=SimuData2(0,1345,aug=1)
ValSet=SimuData2(1345,1495,aug=1) #Numbers specific to Indlist
TrainLoader=DataLoader(TrainSet, batch_size=batch_size,shuffle=True, num_workers=num_workers)
ValLoader=DataLoader(ValSet, batch_size=batch_size,shuffle=True, num_workers=num_workers)
loss_train = []; loss_val = []
for epoch in range(num_epochs):
reg_model.train()
for t, data in enumerate(TrainLoader, 0):
optimizer.zero_grad()
inp=data[0].to(device)
mask = mask_model(inp) > thres
Y_pred = reg_model(inp)
#loss =((Y_pred - target).pow(2)*torch.exp(target*8.)).mean() if mask.any()>0 else (Y_pred - target)*0.
loss = criterion(Y_pred, data[1].to(device),mask)
loss_train.append(loss.item())
loss.backward()
optimizer.step()
np.savetxt(dire+'trainReg.dat',loss_train)
reg_model.eval()
_loss=0
for t_val, data in enumerate(ValLoader,0):
with torch.no_grad():
inp=data[0].to(device)
mask = mask_model(inp) > thres
Y_pred = reg_model(inp)#(reg_model(inp))[mask] #target=(data[1].to(device))[mask]
_loss += criterion(Y_pred, data[1].to(device),mask).item()
loss_val.append(_loss/(t_val+1))
np.savetxt(dire+'valReg.dat',loss_val)
if( _loss/(t_val+1) < best_val):
print('Epoch {}/{}'.format(epoch, num_epochs - 1))
print('-' * 10)
best_val= _loss/(t_val+1)
print('Saving model loss:',best_val)
torch.save(reg_model.state_dict(),dire+'BestReg.pt')
time_elapsed = time.time() - start_time
print('Time {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))
time_elapsed = time.time() - start_time
print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))