def __init__(self, args): """ :param self.train_frames: how many frame are used for training :param self.num_of_vid: number of videos :param self.extractor: pretrained model for feature extraction :param self.G: what I want to training model :param self.optimizerG: optimizer (default: Adam) :param self.train_data: training data loader :param self.device: gpu or cpu device(default: gpu) """ 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 # prtrained 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(shortcut_type='B', cardinality=32, sample_size=self.crop_size, frame_num=self.train_frames, num_of_vid=self.num_of_vid) self.G.train() # optimizer self.optimizerG = optim.Adam(self.G.parameters(), lr=args.lr, betas=(0.9, 0.999), eps=1e-8) # loss self.CE_loss = nn.CrossEntropyLoss() self.L1_loss = nn.L1Loss() self.L2_loss = nn.MSELoss() self.Softmax = nn.Softmax() # data self.train_data = train_data_loader # gpu self.device0 = device0 self.device1 = device1 self.epochs = args.epochs self.avg_G_loss_arr = [] self.checkpoint = args.checkpoint self.checkpoint_set() self.lr_decay = args.lr_decay # cuda if torch.cuda.is_available(): self.extractor.to(self.device0) self.G.cuda(self.device1)
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