Пример #1
0
    def __init__(self, args):
        self.train_frames = args.train_frames
        self.num_of_vid = args.num_of_vid
        self.crop_size = args.crop_size
        self.train_frames = args.train_frames
        self.pretrain_path = pretrain_path
        print("===> Building model")
        # Model Setting
        self.extractor = resnet101(num_classes=400,
                                   shortcut_type='B',
                                   cardinality=32,
                                   sample_size=self.crop_size,
                                   sample_duration=self.train_frames)

        # load pretrained model
        weight = get_pretrain_weight(self.pretrain_path, self.extractor)
        self.extractor.load_state_dict(weight)
        self.extractor.eval()

        self.G = Generator(in_channel=3,
                           crop_size=self.crop_size,
                           frame_num=self.train_frames,
                           num_of_vid=self.num_of_vid)
        self.G.train()
        # Discriminator
        self.D = Context_Discriminator(in_channel=3)
        self.D.train()
        # optimizer
        self.optimizerG = optim.Adam(self.G.parameters(),
                                     lr=args.lr_G,
                                     betas=(0.9, 0.999),
                                     eps=1e-8)
        self.optimizerD = optim.Adam(self.D.parameters(),
                                     lr=args.lr_D,
                                     betas=(0.9, 0.999),
                                     eps=1e-8)
        # loss
        self.BCE_loss = nn.BCELoss()
        self.CE_loss = nn.CrossEntropyLoss()
        # data
        self.train_data = train_data_loader
        # gpu
        self.device0 = device0
        self.device1 = device1

        self.epochs = args.epochs
        self.avg_G_loss_arr = []
        self.avg_D_loss_arr = []
        self.avg_G_label_loss_arr = []
        self.checkpoint = args.checkpoint
        self.checkpoint_set()
        self.lr_decay = args.lr_decay
        self.condi_frame = args.prev_frame

        # cuda
        if torch.cuda.is_available():
            self.extractor.to(self.device0)
            self.G.cuda(self.device1)
            self.D.cuda(self.device1)
Пример #2
0
    def __init__(self, args):
        self.test_frame = args.test_frame
        self.num_of_vid = args.num_of_vid
        self.crop_size = args.crop_size
        self.test_frame = args.test_frame
        print("===> Building model")
        # prtrained model setting
        self.extractor = resnet101(num_classes=400,
                                   shortcut_type='B',
                                   cardinality=32,
                                   sample_size=self.crop_size,
                                   sample_duration=self.test_frame)

        # load pretrained model with eval mode
        weight = get_pretrain_weight(args.pretrained_model_path,
                                     self.extractor)
        self.extractor.load_state_dict(weight)
        self.extractor.eval()

        self.G = torch.load(args.train_model)
        self.G.eval()

        # data loadeer
        self.test_data = test_data_loader

        # cuda GPU devices setting
        self.device0 = device0
        self.device1 = device1
        if torch.cuda.is_available():
            self.extractor.cuda(self.device1)
            self.G.cuda(self.device1)

        self.log_dir = args.log_path
        self.original_vid_dir = args.original_vid_dir
        self.demo_dir = args.demo_dir
        self.out_fps = args.out_fps
        self.height = args.height
        self.width = args.width
        self.in_fps = args.fps
        self.prev_frame = args.prev_frame