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)
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