rnn_basic_classification.py

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader

import sdc.datasets.uci as uci



class TSDataset(Dataset):
  def __init__(self, train=True, raw=True):
        self.is_train = train
        self.train, self.test = uci.load_har(raw=raw)
        
  def __getitem__(self, item):
        if self.is_train:
            data, label = self.train
        else:
            data, label = self.test
        
        return data[item].astype(np.float32), label[item][0].astype(np.int64) - 1
        
  def __len__(self):
        if self.is_train:
            data, label = self.train
        else:
            data, label = self.test
            
        return len(label)
    
class Network(nn.Module):
    def __init__(self, in_size, hidden_size, n_class):
        super(Network, self).__init__()
        self.rnn = nn.LSTM(in_size, hidden_size, batch_first=True)
        self.fc = nn.Linear(hidden_size, n_class)

    def forward(self, x):
        out, state = self.rnn(x)
        out = self.fc(out[:, -1, :])
        return out
    

train_ds = TSDataset(train=True)
test_ds = TSDataset(train=False)
train_loader = DataLoader(train_ds, batch_size=256, shuffle=False)
test_loader = DataLoader(test_ds, batch_size=256, shuffle=False)

model = Network(9, 64, 6)
optimizer = optim.Adam(model.parameters())
criterion = nn.CrossEntropyLoss()


model.train()
for epoch in range(50):
    for x, y in train_loader:
        optimizer.zero_grad()
        output = model(x)
        loss = criterion(output, y)
        loss.backward()
        optimizer.step()
    print(loss)
    


correct = 0.0
total = 0

model.eval()
with torch.no_grad():
    for x, y in test_loader:
        output = model(x)
        loss = criterion(output, y)
        total += y.shape[0]
        _, predicted = torch.max(output, 1)
        correct += (predicted == y).sum().float().cpu().item()

print(correct / total)