CNN.py 2.65 KB
Newer Older
Ard Kastrati's avatar
Ard Kastrati committed
1
# -*- coding: utf-8 -*-
2
import torch
Ard Kastrati's avatar
Ard Kastrati committed
3
4
5
6
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torch.utils.data
okiss's avatar
okiss committed
7
8
from utils.utils import *

Ard Kastrati's avatar
Ard Kastrati committed
9
10
from config import config

Ard Kastrati's avatar
Ard Kastrati committed
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv1d(in_channels=129, out_channels=258, kernel_size=5, stride=1, padding=2)
        self.pool = nn.MaxPool1d(5)
        self.conv2 = nn.Conv1d(in_channels=258, out_channels=64, kernel_size=5, stride=1, padding=2)
        self.fc1 = nn.Linear(64*20, 120)
        self.fc2 = nn.Linear(120, 60)
        self.fc3 = nn.Linear(60, 2)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 64*20)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

31
def run(trainX, trainY):
Ard Kastrati's avatar
Ard Kastrati committed
32
33
34
35
36
37
38
39
40
41
42
43
    #load the data
    dataset = torch.utils.data.TensorDataset(trainX, trainY)
    trainloader = torch.utils.data.DataLoader(dataset, batch_size=2)

    # define the network
    net = Net()

    # define the optimizer
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

    # train
44
    hist = train(trainloader=trainloader, net=net, optimizer=optimizer, criterion=criterion)
Ard Kastrati's avatar
Ard Kastrati committed
45
46

    # save our trained model
47
48
49
    # PATH = '../cifar_net.pth'
    # torch.save(net.state_dict(), PATH)
    # Newly added lines below
Ard Kastrati's avatar
Ard Kastrati committed
50
51
52
	# save_logs(hist, config['model_dir'], config['model'], pytorch=True)
	# save_model_param(config['model_dir'], config['model'], pytorch=True)
	# plot_loss_torch(loss)
Ard Kastrati's avatar
Ard Kastrati committed
53
54

def train(trainloader, net, optimizer, criterion, epoch=50):
okiss's avatar
okiss committed
55
    loss=[]
Ard Kastrati's avatar
Ard Kastrati committed
56
57
    for epoch in range(2):  # loop over the dataset multiple times
        running_loss = 0.0
okiss's avatar
okiss committed
58
        loss_values=[]
Ard Kastrati's avatar
Ard Kastrati committed
59
60
61
62
63
64
65
66
67
68
69
70
        for i, data in enumerate(trainloader, 0):
            # get the inputs; data is a list of [inputs, labels]
            inputs, labels = data

            # zero the parameter gradients
            optimizer.zero_grad()

            # forward + backward + optimize
            outputs = net(inputs)
            loss = criterion(outputs, labels.squeeze(1))
            loss.backward()
            optimizer.step()
okiss's avatar
okiss committed
71
            
Ard Kastrati's avatar
Ard Kastrati committed
72
            # print statistics
okiss's avatar
okiss committed
73
74
75
            run_loss = loss.item()
            running_loss+=run_loss
            loss_values.append(run_loss)
76
            if i % 200 == 0:  # print every 200 mini-batches
Ard Kastrati's avatar
Ard Kastrati committed
77
                print('[%d, %5d] loss: %.3f' %
78
                    (epoch + 1, i + 1, running_loss / 200))
Ard Kastrati's avatar
Ard Kastrati committed
79
                running_loss = 0.0
okiss's avatar
okiss committed
80
81
82
        los=np.mean(loss_values)
        loss.append(los)
    return loss
Ard Kastrati's avatar
Ard Kastrati committed
83
    print('Finished Training')