import torch
import pickle
import pdb
import os
import numpy as np
import matplotlib.pyplot as plt
class MLP(torch.nn.Module):
def __init__(self, in_feats_dim, hid_feats_dim, out_feats_dim):
super().__init__()
self.L1 = torch.nn.Linear(in_feats_dim, hid_feats_dim, True)
self.L2 = torch.nn.Linear(hid_feats_dim, out_feats_dim, True)
def forward(self, inputs):
hids = self.L1(inputs)
hids = torch.relu(hids)
outs = self.L2(hids)
return outs
def data_process(features, labels):
# remove -200
labels = labels.reshape(-1, 1)
datas = np.concatenate((features, labels), axis=1)
nanrows = np.where(datas == -200)[0]
# delete 366 rows
datas = np.delete(datas, nanrows, axis=0)
features, labels = datas[:, :-1], datas[:, -1]
# features normalization
# pdb.set_trace()
mu = np.mean(features, axis=0)
std = np.std(features, axis=0)
assert 0. not in std, "std is zero"
features = (features - mu) / std
# change numpy.ndarry to tensor
features = torch.FloatTensor(features)
labels = torch.FloatTensor(labels)
return features, labels
nfeats, labels = pickle.load(open("data.pkl", 'rb'))
features, labels = nfeats, labels[:, 2]
# data process
# pdb.set_trace()
features, labels = data_process(features, labels)
# divide into trainset and testset
trainN = int(nfeats.shape[0] * 0.75)
train_feats, train_labels = features[:trainN], labels[:trainN]
test_feats, test_labels = features[trainN:], labels[trainN:]
infeatsdim, outfeatsdim = train_feats.shape[1], 1 # train_labels.shape[1]
# super-parameters
hidfeatsdim = 20
lr = 0.01
weight_decay = 1e-5
epochs = 150 # 230-240
batchsize = 64
loadmodel = False
# model and optimizer
model = MLP(infeatsdim, hidfeatsdim, outfeatsdim)
optimizer = torch.optim.Adam(model.parameters(),
lr=lr,
weight_decay=weight_decay)
loss_his = []
# load model
saved_path = "checkpoint"
if loadmodel and os.path.exists(saved_path):
checkpoint = torch.load(saved_path)
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
# loss function
loss_fn = torch.nn.MSELoss(reduction='mean') #
# train
for i in range(epochs):
model.train()
optimizer.zero_grad()
outputs = model(train_feats)
train_loss = loss_fn(outputs, train_labels)
train_loss.backward()
optimizer.step()
print(f"Epoch{i}: trian loss {train_loss:.3f}", end=" ")
loss_his.append(float(train_loss))
# test
model.eval()
outputs = model(test_feats)
test_loss = loss_fn(outputs, test_labels)
print(f"test loss {test_loss:.3f}")
# save model
if loadmodel:
torch.save({
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict()
}, saved_path)