dennistrujillo commited on
Commit
c59d416
1 Parent(s): bbf4478

added inference.py

Browse files
Files changed (1) hide show
  1. inference.py +41 -0
inference.py ADDED
@@ -0,0 +1,41 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from model import model_init, BraggNN
2
+ from torch.utils.data import DataLoader
3
+ from dataset import PatchWiseDataset
4
+ import torch, argparse, os
5
+ from util import str2bool, str2tuple, s2ituple
6
+ from plot import plot_loss, plot_error
7
+ import numpy as np
8
+
9
+ parser = argparse.ArgumentParser(description='Bragg peak finding for HEDM.')
10
+ parser.add_argument('-p_file', type=str, default="debug", help='frame/patch metadata h5 file')
11
+ parser.add_argument('-f_file', type=str, default="debug", help='frame/patch h5 file')
12
+ parser.add_argument('-m_file', type=str, help='model state_dict file')
13
+ parser.add_argument('-psz', type=int, default=15, help='working patch size')
14
+ parser.add_argument('-fcsz', type=s2ituple, default='16_8_4_2', help='size of dense layers')
15
+ parser.add_argument('-expName',type=str, default="debug", help='Experiment name')
16
+ args, unparsed = parser.parse_known_args()
17
+
18
+ def main(args):
19
+ device = torch.device("cpu")#"cuda" if torch.cuda.is_available() else "cpu")
20
+ #model = BraggNN(imgsz=args.psz, fcsz=args.fcsz)
21
+ if device.type == 'cpu':
22
+ model=torch.load(args.m_file,map_location=torch.device('cpu'))
23
+ else:
24
+ model=torch.load(args.m_file)
25
+ model.eval()
26
+
27
+ ds_valid = PatchWiseDataset(psz=args.psz, rnd_shift=0, use='validation', pfile=args.p_file, ffile=args.f_file)
28
+ dl_valid = DataLoader(dataset=ds_valid, batch_size=len(ds_valid), shuffle=True, \
29
+ num_workers=4, drop_last=False, pin_memory=True)
30
+ mse=0.0
31
+ with torch.no_grad():
32
+ for i, (inputs, labels) in enumerate(dl_valid):
33
+ inputs = inputs.to(device)
34
+ y_pred = model(inputs)
35
+ y_pred = y_pred.cpu().numpy()
36
+ labels = labels.cpu().numpy()
37
+ plot_error(y_pred,labels,args.expName)
38
+ mse += np.power(y_pred[:,0] - labels[:,0],2) + np.power(y_pred[:,1] - labels[:,1],2)
39
+ if __name__ == "__main__":
40
+ main(args)
41
+