Exemple #1
0
class NN_Trainer(object):
    def __init__(self, **kwargs):
        self.batch_size = kwargs['batch_size']
        self.lr = kwargs['learning_rate']
        self.max_epochs = kwargs['max_epochs']
        self.momentum = kwargs['momentum']
        self.network_config = kwargs['network']

    def build_model(self):
        # build network
        if self.network_config == "LeNet":
            self.network = LeNet()
        elif self.network_config == "ResNet":
            #self.network=ResNet18()
            self.network = ResNetSplit18(1)
        # set up optimizer
        self.optimizer = torch.optim.SGD(self.network.parameters(),
                                         lr=self.lr,
                                         momentum=self.momentum)
        self.criterion = torch.nn.CrossEntropyLoss()

    def train_and_validate(self, train_loader, test_loader):
        # iterate of epochs
        for i in range(self.max_epochs):
            # change back to training mode
            self.network.train()
            for batch_idx, (data, y_batch) in enumerate(train_loader):
                iter_start_time = time.time()
                data, target = Variable(data), Variable(y_batch)
                self.optimizer.zero_grad()
                ################# backward on normal model ############################
                '''
                logits = self.network(data)
                loss = self.criterion(logits, target)
                '''
                #######################################################################

                ################ backward on splitted model ###########################
                logits = self.network(data)
                logits_1 = Variable(logits.data, requires_grad=True)
                loss = self.criterion(logits_1, target)
                loss.backward()
                self.network.backward_single(logits_1.grad)
                #######################################################################
                tmp_time_0 = time.time()

                for param in self.network.parameters():
                    # get gradient from layers here
                    # in this version we fetch weights at once
                    # remember to change type here, which is essential
                    grads = param.grad.data.numpy().astype(np.float64)

                duration_backward = time.time() - tmp_time_0

                tmp_time_1 = time.time()
                self.optimizer.step()
                duration_update = time.time() - tmp_time_1

                # calculate training accuracy
                prec1, prec5 = accuracy(logits.data, y_batch, topk=(1, 5))
                # load the training info
                print(
                    'Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}  Prec@1: {}  Prec@5: {}  Time Cost: {}'
                    .format(i, batch_idx * len(data),
                            len(train_loader.dataset),
                            100. * batch_idx / len(train_loader), loss.data[0],
                            prec1.numpy()[0],
                            prec5.numpy()[0],
                            time.time() - iter_start_time))
            # we evaluate the model performance on end of each epoch
            self.validate(test_loader)

    def validate(self, test_loader):
        self.network.eval()
        test_loss = 0
        correct = 0
        prec1_counter_ = prec5_counter_ = batch_counter_ = 0
        for data, y_batch in test_loader:
            data, target = Variable(data, volatile=True), Variable(y_batch)
            output = self.network(data)
            test_loss += F.nll_loss(
                output, target,
                size_average=False).data[0]  # sum up batch loss
            #pred = output.data.max(1, keepdim=True)[1] # get the index of the max log-probability
            #correct += pred.eq(target.data.view_as(pred)).cpu().sum()
            prec1_tmp, prec5_tmp = accuracy(output.data, y_batch, topk=(1, 5))
            prec1_counter_ += prec1_tmp.numpy()[0]
            prec5_counter_ += prec5_tmp.numpy()[0]
            batch_counter_ += 1
        prec1 = prec1_counter_ / batch_counter_
        prec5 = prec5_counter_ / batch_counter_
        test_loss /= len(test_loader.dataset)
        print('Test set: Average loss: {:.4f}, Prec@1: {} Prec@5: {}'.format(
            test_loss, prec1, prec5))
Exemple #2
0
class NN_Trainer(object):
    def __init__(self, **kwargs):
        self.batch_size = kwargs['batch_size']
        self.lr = kwargs['learning_rate']
        self.max_epochs = kwargs['max_epochs']
        self.momentum = kwargs['momentum']
        self.network_config = kwargs['network']

    def build_model(self):
        # build network
        if self.network_config == "LeNet":
            self.network = LeNet()
        elif self.network_config == "ResNet":
            self.network = ResNet18()
        # set up optimizer
        self.optimizer = torch.optim.SGD(self.network.parameters(),
                                         lr=self.lr,
                                         momentum=self.momentum)
        self.criterion = torch.nn.CrossEntropyLoss()

    def train(self, train_loader):
        self.network.train()

        # iterate of epochs
        for i in range(self.max_epochs):
            for batch_idx, (data, y_batch) in enumerate(train_loader):
                iter_start_time = time.time()
                data, target = Variable(data), Variable(y_batch)
                self.optimizer.zero_grad()
                logits = self.network(data)
                loss = self.criterion(logits, target)
                tmp_time_0 = time.time()
                loss.backward()

                #for key_name, param in self.network.state_dict().items():
                #    print(param)
                #    print("----------------------------------------------------------------")
                #exit()
                for param in self.network.parameters():
                    # get gradient from layers here
                    # in this version we fetch weights at once
                    # remember to change type here, which is essential
                    grads = param.grad.data.numpy().astype(np.float64)

                duration_backward = time.time() - tmp_time_0

                tmp_time_1 = time.time()
                self.optimizer.step()
                duration_update = time.time() - tmp_time_1

                print("backward duration: {}".format(duration_backward))
                print("update duration: {}".format(duration_update))
                # calculate training accuracy
                prec1, prec5 = accuracy(logits.data, y_batch, topk=(1, 5))
                # load the training info
                print(
                    'Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}  Prec@1: {}  Prec@5: {}  Time Cost: {}'
                    .format(i, batch_idx * len(data),
                            len(train_loader.dataset),
                            100. * batch_idx / len(train_loader), loss.data[0],
                            prec1.numpy()[0],
                            prec5.numpy()[0],
                            time.time() - iter_start_time))