Upload New File

Aufgabe 2: Spectrometer Alignment/Spectrometer_Alignment.py

+import h5py
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from h5py import Dataset
+from sklearn.preprocessing import MinMaxScaler
+from torch.utils.data import DataLoader
+from torch.optim import SGD, Adam
+from torch.nn import Linear, Sigmoid, Module, MSELoss, BCELoss
+from tqdm import tqdm
+
+# f hat 100k keys, jeder key hat eine hdf5 group als value
+# jede group hat 2 keys. in key 'X' ist ein 2d numpy array
+# welches die 200x200 pixel des bildes speichert (dset[70:71,:]
+# zugriff auf zeile 71) und in key'Y' sind die drei
+# korrespondierenden achsenwerte der kamera X, Y und Z
+
+
+def prepare_data(path):
+    # load h5 file
+    sims = 3000  # only half the data
+    x = []
+    y = []
+    with h5py.File(path, 'r') as f:
+        for key in range(sims):
+            X = np.array(f['/' + str(key) + '/X'])
+            Y = np.array(f['/' + str(key) + '/Y'])
+            X = MinMaxScaler().fit_transform(X)  # normalise the data between 0 and 1
+            # Y = MinMaxScaler().fit_transform(Y)
+            # print(Y[:])
+            if np.amax(X) > 0:  # ignore the blank images (no data)
+                x.append(X)
+                y.append(Y)
+    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+    print(device)
+    x = torch.tensor(x, dtype=torch.float32, device=device)  # input
+    y = torch.tensor(y, dtype=torch.float32, device=device)  # output
+    print('x.shape', np.shape(x))
+    print('y.shape', np.shape(y))
+    print(x.element_size() * x.nelement())
+
+    # split test and train data
+    data_length = np.shape(x)[0]
+    train_length = data_length - 1000
+    train_x = x[0:train_length]
+    test_x = x[train_length:data_length]
+    train_y = y[0:train_length]
+    test_y = y[train_length:data_length]
+    train = Data(train_x, train_y)
+    test = Data(test_x, test_y)
+    return train, test
+
+    # train, test = dataset.get_splits()
+    # trainloader = torch.utils.data.DataLoader(x, batch_size=8, shuffle=True, num_workers=0)
+
+
+def train_model(train_dl, model, lr, numEpochs, optimizerFlag, modelSaveFileName, criterionFlag, startEpoch):
+    if optimizerFlag == 'A':
+        optimizer = Adam(model.parameters(), lr=lr, betas=(0.9, 0.999), eps=1e-08, weight_decay=0, amsgrad=False)
+    else:
+        optimizer = SGD(model.parameters(), lr=lr, momentum=0.9)
+    if criterionFlag == 'MSE':
+        criterion = MSELoss()
+    else:
+        criterion = BCELoss()
+    for epoch in range(startEpoch + 1, numEpochs + 1):
+        # enumerate mini batches
+        loop = tqdm(enumerate(train_dl), total=len(train_dl), leave=False)
+        for i, (inputs, targets) in loop:
+            # clear the gradients
+            optimizer.zero_grad()
+            # compute the model output
+            yhat = model(inputs)
+            # calculate loss
+            loss = criterion(yhat, targets)
+            # credit assignment
+            loss.backward()
+            # update model weights
+            optimizer.step()
+            # upgrade progress bar
+            loop.set_description(f"Epoch [{epoch}/{numEpochs}]")
+            # loop.set_postfix(loss=loss.item(),acc=torch.rand(1).item())
+        if epoch % 10 == 0:
+            torch.save({
+                'epoch': epoch,
+                'model_state_dict': model.state_dict(),
+                'optimizer_state_dict': optimizer.state_dict(),
+                'lr': lr,
+            }, modelSaveFileName + str(epoch) + ".pt")
+
+
+# TODO anpassen an unsere Situation
+class Net(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = nn.Conv2d(1, 8, kernel_size=5)  # 1d pictures, 8 pictures at a time, 5x5 kernel size
+        self.pool = nn.MaxPool2d(2, 2)
+        self.conv2 = nn.Conv2d(1, 8, kernel_size=5)
+        self.fc1 = nn.Linear(200, 120)
+        self.fc2 = nn.Linear(120, 84)
+        self.fc3 = nn.Linear(84, 10)
+
+    def forward(self, x):
+        x = self.pool(F.relu(self.conv1(x)))
+        x = self.pool(F.relu(self.conv2(x)))
+        x = torch.flatten(x, 1)  # flatten all dimensions except batch
+        x = F.relu(self.fc1(x))
+        x = F.relu(self.fc2(x))
+        x = self.fc3(x)
+        return x
+
+
+class Data(Dataset):
+    def __init__(self, x, y):
+        self.x = x
+        self.y = y
+        self.len = len(x)
+
+    def __getitem__(self, index):
+        self.xc = self.x[index]
+        self.yc = self.y[index]
+        return self.xc, self.yc
+
+    def __len__(self):
+        return self.len
+
+
+path = "training_data_alignment.h5"
+batch_size = 8
+num_Epochs = 100
+modelSaveFileName = 'model.pt'
+optimizerFlag = 'A'
+criterionFlag = 'MSE' # MSE oder BCE
+startEpoch = 1
+learning_rate = 0.01
+
+train, test = prepare_data(path)
+train_dl = DataLoader(train, batch_size=batch_size, shuffle=True, num_workers=0, pin_memory=False)
+test_dl = DataLoader(test, batch_size=batch_size, shuffle=True, num_workers=0, pin_memory=False)
+net = Net()
+
+train_model(train_dl, net, learning_rate, num_Epochs, optimizerFlag, modelSaveFileName, criterionFlag, startEpoch)
+
+# print(net)