chirp/train_classifer.py at master · kunglab/chirp · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
import numpy as np
import os, sys
import argparse

import chainer
import chainer.functions as F
import chainer.links as L
from chainer import training
from common.dataset import Dataset, LabeledDataset
from chainer.training import extensions
from chainer.datasets import get_mnist
from chainer import serializers
from sklearn import metrics

from common.net import Alex, VTCNN2
import cupy

from common import dataset


## Alex net ##


model_map = {"AlexStock": Alex, "VTCNN2" : VTCNN2}

parser = argparse.ArgumentParser(description='Chainer example: MNIST')
parser.add_argument('--batchsize', '-b', type=int, default=100,
                    help='Number of images in each mini-batch')
parser.add_argument('--epoch', '-e', type=int, default=20,
                    help='Number of sweeps over the dataset to train')
parser.add_argument('--gpu', '-g', type=int, default=-1,
                    help='GPU ID (negative value indicates CPU)')
parser.add_argument('--out', '-o', default='result_classifier',
                    help='Directory to output the result')
parser.add_argument('--unit', '-u', type=int, default=1000,
                    help='Number of units')
parser.add_argument('--model_type', '-t', type=str, default="AlexStock",
                    help='Which Model to run (AlexStock, VTCNN2)')
args = parser.parse_args()


if not os.path.exists(args.out):
    os.makedirs(args.out)

###### SETUP DATASET #####
noise_levels = range(-18, 20, 2)
RFdata_train = dataset.RFModLabeled(noise_levels=noise_levels, test=False)
RFdata_test = dataset.RFModLabeled(noise_levels=noise_levels, test=True)

num_classes = np.unique(RFdata_train.ys).shape[0]

if False:
    xs = RFdata_train.xs[:,:,0,:] + RFdata_train.xs[:,:,1,:]
    ys = RFdata_train.xs[:,:,0,:] - RFdata_train.xs[:,:,1,:]
    RFdata_train.xs[:,:,0,:] = xs
    RFdata_train.xs[:,:,1,:] = ys

    xs = RFdata_test.xs[:,:,0,:] + RFdata_test.xs[:,:,1,:]
    ys = RFdata_test.xs[:,:,0,:] - RFdata_test.xs[:,:,1,:]
    RFdata_test.xs[:,:,0,:] = xs
    RFdata_test.xs[:,:,1,:] = ys

# train model
if args.model_type == "AlexStock":
    model = L.Classifier(model_map[args.model_type](num_classes, init_weights=True, filter_height=2))
else:
    model = L.Classifier(model_map[args.model_type](num_classes))

if args.gpu >= 0:
	chainer.cuda.get_device_from_id(args.gpu).use()
	# model.to_gpu(args.gpu)


#optimizer = chainer.optimizers.Adam(alpha=0.0001, beta1=0.0, beta2=.9)
optimizer = chainer.optimizers.Adam()
optimizer.setup(model)
train_iter = chainer.iterators.SerialIterator(RFdata_train, args.batchsize)
test_iter = chainer.iterators.SerialIterator(RFdata_test, args.batchsize,
                                             repeat=False, shuffle=False)
# updater = training.StandardUpdater(train_iter, optimizer, device=args.gpu)
updater = training.ParallelUpdater(
        train_iter,
        optimizer,
        # The device of the name 'main' is used as a "master", while others are
        # used as slaves. Names other than 'main' are arbitrary.
        devices={'main': 0, 'second': 1},
    )
trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=args.out)
trainer.extend(extensions.Evaluator(test_iter, model, device=args.gpu))
trainer.extend(extensions.LogReport())
trainer.extend(extensions.PrintReport(
    ['epoch', 'main/loss', 'validation/main/loss',
     'main/accuracy', 'validation/main/accuracy', 'elapsed_time']))
trainer.extend(extensions.ProgressBar())
trainer.run()

serializers.save_npz(os.path.join(args.out, 'main_classifer_greaterthan10_regularized.npz'), model)

x, y = RFdata_test.xs, RFdata_test.ys
xp = np if args.gpu < 0 else cupy

pred_ys = xp.zeros(y.shape)


chainer.config.train = False
for i in range(0, len(x), args.batchsize):
    x_batch = xp.array(x[i:i + args.batchsize])
    y_batch = xp.array(y[i:i + args.batchsize])
    y_pred = model.predictor(x_batch)
    acc = model.accfun(y_pred, y_batch)
    acc = chainer.cuda.to_cpu(acc.data)
    # print "Accuracy: ", acc
    pred_ys[i:i + args.batchsize] = np.argmax(y_pred._data[0], axis=1)
chainer.config.train = True


np.savez(os.path.join(args.out,'pred_ys__main_classifer_greaterthan10_regularized.npz'), pred_ys = chainer.cuda.to_cpu(pred_ys))

cm = metrics.confusion_matrix(chainer.cuda.to_cpu(y), chainer.cuda.to_cpu(pred_ys))
print cm

cor = np.sum(np.diag(cm))
ncor = np.sum(cm) - cor
print "Overall Accuracy: ", cor / float(cor+ncor)