예제 #1
0
    def __init__(self, root, mode, test_dir, train_dir, save_model_withname=None,\
                 save_error_withname=None, checkpoint=None):
        self.root = root
        self.mode = mode
        self.test_dir = test_dir
        self.train_dir = train_dir
        self.save_model_withname = save_model_withname
        self.save_error_withname = save_error_withname
        self.checkpoint = checkpoint
        self.batch_size = 50
        self.learning_rate = 0.0001
        self.validation_loop = 0

        if(self.mode=='train'):
            self.writer = tensorboardX.SummaryWriter(comment="train")
        else:
            self.writer = tensorboardX.SummaryWriter(comment="test")
        # setup dataset
        self.train_transforms = transforms.Compose([videotransforms.RandomCrop(112),
                                           videotransforms.RandomHorizontalFlip(),])
        self.test_transforms = transforms.Compose([videotransforms.CenterCrop(112)])

        self.dataset = VisualTactile(self.root, self.train_dir, self.train_transforms)
        self.dataloader = torch.utils.data.DataLoader(self.dataset, batch_size=self.batch_size, shuffle=True, num_workers=1, pin_memory=True)

        self.val_dataset = VisualTactile(self.root, self.test_dir, self.test_transforms)
        self.val_dataloader = torch.utils.data.DataLoader(self.val_dataset, batch_size=1, shuffle=False, num_workers=1, pin_memory=True)

#         self.dataloaders = {'train': self.dataloader, 'val': self.val_dataloader}
#         self.datasets = {'train': self.dataset, 'val': self.val_dataset}

        self.model, self.optimizer, self.scheduler = self.load_model(self.checkpoint)
class Run_Model(object):

    def __init__(self, root, mode, test_dir, train_dir, save_model_withname=None,\
                 save_error_withname=None, checkpoint=None):
        self.root = root
        self.mode = mode
        self.test_dir = test_dir
        self.train_dir = train_dir
        self.save_model_withname = save_model_withname
        self.save_error_withname = save_error_withname
        self.checkpoint = checkpoint
        self.batch_size = 50
        self.learning_rate = 0.0001

        if (self.mode == 'train'):
            self.writer = tensorboardX.SummaryWriter(comment="train")
        else:
            self.writer = tensorboardX.SummaryWriter(comment="test")
        # setup dataset
        self.train_transforms = transforms.Compose([
            videotransforms.RandomCrop(112),
            videotransforms.RandomHorizontalFlip(),
        ])
        self.test_transforms = transforms.Compose(
            [videotransforms.CenterCrop(112)])

        self.dataset = VisualTactile(self.root, self.train_dir,
                                     self.train_transforms)
        self.dataloader = torch.utils.data.DataLoader(
            self.dataset,
            batch_size=self.batch_size,
            shuffle=True,
            num_workers=2,
            pin_memory=True)

        self.val_dataset = VisualTactile(self.root, self.test_dir,
                                         self.test_transforms)
        self.val_dataloader = torch.utils.data.DataLoader(self.val_dataset,
                                                          batch_size=1,
                                                          shuffle=False,
                                                          num_workers=1,
                                                          pin_memory=True)

        #         self.dataloaders = {'train': self.dataloader, 'val': self.val_dataloader}
        #         self.datasets = {'train': self.dataset, 'val': self.val_dataset}

        self.model, self.optimizer, self.scheduler = self.load_model(
            self.checkpoint)

    def load_model(self, checkpoint):
        sm = resnet.resnet18(sample_size=112,
                             sample_duration=18,
                             num_classes=400,
                             shortcut_type='A')
        sm = nn.DataParallel(sm)
        if torch.cuda.is_available():
            pretrain = torch.load("../models/resnet-18-kinetics.pth")
            sm.load_state_dict(pretrain['state_dict'])
        else:
            pretrain = torch.load("../models/resnet-18-kinetics.pth",
                                  map_location="cpu")
            #            pretrain = torch.load("../../../out/resnet/resnet-18-kinetics.pth", map_location="cpu")
            sm.load_state_dict(pretrain['state_dict'])
        sm = self.freeze_network_layer(sm)

        net = FusionNet(sm)
        if torch.cuda.is_available():
            net.cuda()
            # net = nn.DataParallel(net)
        optimizer = optim.Adam(net.parameters(),
                               lr=self.learning_rate,
                               weight_decay=0.0000001)
        lr_sched = optim.lr_scheduler.MultiStepLR(optimizer, [10, 20, 25])

        #checkpoint in case of training
        if (checkpoint is not None):
            if torch.cuda.is_available():
                data = torch.load(checkpoint)
            else:
                data = torch.load(checkpoint,
                                  map_location=lambda storage, loc: storage)
            net.load_state_dict(data['model_state'])
            optimizer.load_state_dict(data['optimizer_state'])
            lr_sched.load_state_dict(data['scheduler_state'])
        return net, optimizer, lr_sched

    def freeze_network_layer(self, model):
        for para in model.parameters():
            para.requires_grad = False
        return model

    def save_checkpoint(self, model, optimizer, scheduler, epoch, save_model):
        data = {
            'model_state': model.state_dict(),
            'optimizer_state': optimizer.state_dict(),
            'scheduler_state': scheduler.state_dict(),
            'epoch': epoch + 1
        }
        torch.save(data, 'model_%d' % (epoch + 1) + save_model + '.tar')

    def train(self):
        with open(self.save_error_withname, 'w') as file:
            file.write("train loss file\n")
        epoch_num = 30
        self.model.train(True)
        for epoch in range(epoch_num):
            print('Step {}/{}'.format(epoch, epoch_num))
            print('-' * 10)
            epoch_time = time.time()

            self.optimizer.zero_grad()
            total_loss = 0
            for i, data in enumerate(self.dataloader):
                vid, lab, path = data

                if torch.cuda.is_available():
                    video = Variable(vid.cuda())
                    label = Variable(lab.cuda())
                else:
                    video = Variable(vid)
                    label = Variable(lab)

                out = self.model(video.float())
                out = out.squeeze(1)
                loss = F.binary_cross_entropy_with_logits(
                    out.float(), label.float())
                total_loss += loss.item()
                loss.backward()
                print('{} Loss: {:.4f} and lr: {}'.format(
                    self.mode, total_loss / (i + 1),
                    self.scheduler.get_lr()[0]))
                with open(self.save_error_withname, 'a') as file:
                    file.write("epoch: {}, Loss: {}, LR: {}\n".format(
                        epoch, total_loss / (i + 1),
                        self.scheduler.get_lr()[0]))

                self.optimizer.step()
                self.optimizer.zero_grad()
                self.writer.add_scalar("error/{}".format(epoch),
                                       total_loss / (i + 1), i)
            self.writer.add_scalar("errorPerEpoch/", total_loss / (i + 1),
                                   epoch)
            self.scheduler.step()
            print("epoch {} :: time {}".format(epoch,
                                               time.time() - epoch_time))
            if ((epoch + 1) % 30 == 0):
                self.save_checkpoint(self.model, self.optimizer,
                                     self.scheduler, epoch,
                                     self.save_model_withname)

    def test(self):
        with open(self.save_error_withname, 'w') as file:
            file.write("test loss file\n")
        self.model.train(False)
        test_TP, test_TN, test_FP, test_FN = 0, 0, 0, 0
        actual_out, predicted_out = 0, 0
        video_num = self.val_dataset.get_num_videos()
        print(video_num, self.val_dataset.get_num_clips())
        b
        count = 0
        for index in range(video_num):
            data = self.val_dataset.get_video_frames(index)
            packed_data, vid_path = data
            print("directory: {}".format(vid_path))
            # iterating though mini clips in a video
            for i, dota in enumerate(packed_data):
                video, label = dota
                label = label.unsqueeze(
                    0)  #because without this its shape is empty
                if torch.cuda.is_available():
                    video = Variable(video.cuda())
                    label = Variable(label.cuda())
                else:
                    video = Variable(video)
                    label = Variable(label)
                video = video.unsqueeze(0)
                out = self.model(video.float())
                print(out.shape, label)
                out = out.squeeze(1)
                print(out.shape, type(label))
                loss = F.binary_cross_entropy_with_logits(
                    out.float(), label.float())
                actual_out = label[0]
                predicted_out = 1 if out[0] > 0 else 0
                print(
                    '{}: Loss: {:.5f} and lr: {:.5f}.... Network output: {:.5f} and actual label: {} '
                    .format(i, loss.item(),
                            self.scheduler.get_lr()[0], out[0], label[0]))
                with open(self.save_error_withname, 'a') as file:
                    file.write(
                        '{}: Loss: {:.5f} and lr: {:.5f}.... Network output: {:.5f} and actual label: {} \n'
                        .format(i, loss.item(),
                                self.scheduler.get_lr()[0], out[0], label[0]))
                self.writer.add_scalar('inference_error/{}'.format(index),
                                       loss.item(), i)
                self.writer.add_scalar('combined_inference_error/',
                                       loss.item(), count)
                count += 1

                if actual_out == predicted_out:
                    if actual_out:
                        test_TP += 1
                    else:
                        test_TN += 1
                else:
                    if actual_out:
                        test_FN += 1
                    else:
                        test_FP += 1
        with open(self.save_error_withname, 'a') as file:
            file.write("-" * 100)
            file.write("\n")
            file.write("Network information\n")
            file.write(
                "learning rate: {}, batch size: {}, saved model name: {}, optimizer : Adam, learning steps: [10,20,25] \n"
                .format(self.learning_rate, self.batch_size,
                        self.save_model_withname))
            file.write("-" * 100)
            file.write("\n")
            file.write("Confusion matrix for test data \n")
            file.write("TP: {}, TN: {}, FP: {}, FN: {} \n".format(
                test_TP, test_TN, test_FP, test_FN))