def __init__(self,
n_classes=10,
latent_sz=32,
ngf=32,
init_type='orthogonal',
init_gain=0.1,
img_sz=32):
super(CGN, self).__init__()
# params
self.batch_size = 1 # default: sample a single image
self.n_classes = n_classes
self.latent_sz = latent_sz
self.label_emb = nn.Embedding(n_classes, n_classes)
init_sz = img_sz // 4
inp_dim = self.latent_sz + self.n_classes
# models
self.f_shape = nn.Sequential(*shape_layers(inp_dim, 1, ngf, init_sz))
self.f_text1 = nn.Sequential(*texture_layers(inp_dim, 3, ngf, init_sz),
nn.Tanh())
self.f_text2 = nn.Sequential(*texture_layers(inp_dim, 3, ngf, init_sz),
nn.Tanh())
self.shuffler = nn.Sequential(Patch2Image(img_sz, 2),
RandomCrop(img_sz))
init_net(self, init_type=init_type, init_gain=init_gain)
def build_basic_network(in_channels, in_size, out_dim, noisy, sigma0, net_file):
conv = _build_default_conv(in_channels)
in_shape = (1, in_channels, in_size, in_size)
fc_in = utils.count_output_size(in_shape, conv)
fc_hid = 512
dims = [fc_in, fc_hid, out_dim]
if noisy:
fc = _build_noisy_fc(dims, sigma0)
else:
fc = _build_fc(dims)
net = BasicNetwork(conv, fc)
utils.init_net(net, net_file)
return net
def build_basic_network(in_channels, in_size, out_dim, noisy, sigma0,
net_file):
conv = _build_default_conv(in_channels)
in_shape = (1, in_channels, in_size, in_size)
fc_in = utils.count_output_size(in_shape, conv)
fc_hid = 512
dims = [fc_in, fc_hid, out_dim]
if noisy:
fc = _build_noisy_fc(dims, sigma0)
else:
fc = _build_fc(dims)
net = BasicNetwork(conv, fc)
utils.init_net(net, net_file)
return net
def build_dueling_network(in_channels, in_size, out_dim, noisy, sigma0, net_file):
conv = _build_default_conv(in_channels)
in_shape = (1, in_channels, in_size, in_size)
fc_in = utils.count_output_size(in_shape, conv)
fc_hid = 512
adv_dims = [fc_in, fc_hid, out_dim]
val_dims = [fc_in, fc_hid, 1]
if noisy:
adv = _build_noisy_fc(adv_dims, sigma0)
val = _build_noisy_fc(val_dims, sigma0)
else:
adv = _build_fc(adv_dims)
val = _build_fc(val_dims)
net = DuelingNetwork(conv, adv, val)
utils.init_net(net, net_file)
return net
def build_dueling_network(in_channels, in_size, out_dim, noisy, sigma0,
net_file):
conv = _build_default_conv(in_channels)
in_shape = (1, in_channels, in_size, in_size)
fc_in = utils.count_output_size(in_shape, conv)
fc_hid = 512
adv_dims = [fc_in, fc_hid, out_dim]
val_dims = [fc_in, fc_hid, 1]
if noisy:
adv = _build_noisy_fc(adv_dims, sigma0)
val = _build_noisy_fc(val_dims, sigma0)
else:
adv = _build_fc(adv_dims)
val = _build_fc(val_dims)
net = DuelingNetwork(conv, adv, val)
utils.init_net(net, net_file)
return net
def __init__(self, n_classes, ndf):
super(DiscConv, self).__init__()
cin = 4 # RGB + Embedding
self.label_embedding = nn.Embedding(n_classes, 1)
def block(cin, cout, ks, st):
return[
nn.Conv2d(cin, cout, ks, stride=st, padding=0, bias=False),
nn.BatchNorm2d(cout),
nn.LeakyReLU(0.2, True),
]
self.model = nn.Sequential(
*block(cin, ndf, 3, 1),
*block(ndf, ndf*2, 3, 1),
*block(ndf*2, ndf*4, 4, 2),
*block(ndf*4, ndf*4, 4, 2),
nn.AvgPool2d(3),
nn.Conv2d(ndf*4, 1, kernel_size=1, stride=1, padding=0, bias=False),
)
init_net(self)
def setup(self):
if self.args.mode == 'train':
self.set_random_seed(self.args.seed)
# 以{'网络名称'：网络对象}的字典的形式申明网络结构（有几个就添加几个）
self.networks = {'LstmPuncNet': LstmPunctuator(self.args.vocab_size, self.args.embedding_dim, self.args.hidden_size,
self.args.num_layers, bidirectional=True, num_class=self.args.num_classes)}
# 网络结构的初始化
self.networks = init_net(self.networks, self.args.init_type, self.args.gpu_ids)
# 优化器
self.optimizers = {'optimizer': torch.optim.SGD(self.networks['LstmPuncNet'].parameters(), lr=self.args.lr, weight_decay=1e-4)}
# 学习率衰减策略
self.schedulers = [get_scheduler(optimizer, self.args) for optimizer in list(self.optimizers.values())]
# 损失函数，可以在这里写一些备用的，方便修改
self.objectives = {'CrossEntropyLoss': torch.nn.CrossEntropyLoss(ignore_index=self.args.ignore_index).to(self.device),
'NLLLoss': torch.nn.NLLLoss().to(self.device)}
else:
self.networks = {'LstmPuncNet': LstmPunctuator(self.args.vocab_size, self.args.embedding_dim, self.args.hidden_size,
self.args.num_layers, bidirectional=True, num_class=self.args.num_classes)}
self.objectives = {'CrossEntropyLoss': torch.nn.CrossEntropyLoss(ignore_index=self.args.ignore_index).to(self.device),
'NLLLoss': torch.nn.NLLLoss().to(self.device)}
self.load_networks(self.networks, self.args.load_epoch)
network_downsampling=network_downsampling,
inlier_percentage=inlier_percentage,
use_store_data=True,
store_data_root=training_data_root,
phase="validation", is_hsv=is_hsv,
num_pre_workers=num_workers, visible_interval=30, rgb_mode="rgb")
train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True,
num_workers=num_workers)
validation_loader = torch.utils.data.DataLoader(dataset=validation_dataset, batch_size=batch_size, shuffle=False,
num_workers=batch_size)
depth_estimation_model_student = models.FCDenseNet57(n_classes=1)
# Initialize the depth estimation network with Kaiming He initialization
depth_estimation_model_student = utils.init_net(depth_estimation_model_student, type="kaiming", mode="fan_in",
activation_mode="relu",
distribution="normal")
# Multi-GPU running
depth_estimation_model_student = torch.nn.DataParallel(depth_estimation_model_student)
# Summary network architecture
if display_architecture:
torchsummary.summary(depth_estimation_model_student, input_size=(3, height, width))
# Optimizer
optimizer = torch.optim.SGD(depth_estimation_model_student.parameters(), lr=max_lr, momentum=0.9)
lr_scheduler = scheduler.CyclicLR(optimizer, base_lr=min_lr, max_lr=max_lr, step_size=num_iter)
# Custom layers
depth_scaling_layer = models.DepthScalingLayer(epsilon=depth_scaling_epsilon)
depth_warping_layer = models.DepthWarpingLayer(epsilon=depth_warping_epsilon)
flow_from_depth_layer = models.FlowfromDepthLayer()
# Loss functions
netG = nets.UNet11(num_classes=1, pretrained='vgg')
# # ngf = 32
# use_dropout = False
# norm_layer = utils.get_norm_layer(norm_type='batch')
# # netG = models.ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=9)
# netG = models.LeakyResnetGenerator(input_nc, output_nc, ngf=6, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=9)
# # netG = models.LeakyResnetGenerator(input_nc, output_nc, ngf=64, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=9)
# ## Unet Generator
# # netG = models.UnetGenerator(input_nc, output_nc, 7, ngf=3, norm_layer=norm_layer, use_dropout=use_dropout)
# utils.init_net(netG, init_type='normal', init_gain=0.02)
summary(netG, input_size=(3, img_size, img_size))
## Building Discriminator
netD = models.Discriminator(input_nc=1, img_size=img_size)
utils.init_net(netD, init_type='normal', init_gain=0.02)
summary(netD, input_size=(1, img_size, img_size))
lr = 0.00002
G_optimizer = Adam(netG.parameters(), lr=lr, betas=(0.9, 0.999))
D_optimizer = Adam(netD.parameters(), lr=lr, betas=(0.9, 0.999))
G_model_path = root / 'G_model.pt'
D_model_path = root / 'D_model.pt'
if G_model_path.exists() and D_model_path.exists():
state = torch.load(str(G_model_path))
netG.load_state_dict(state['model'])
state = torch.load(str(D_model_path))
epoch = state['epoch']
step = state['step']
def __init__(self, args=args):
super().__init__()
self.args = args
# random_seed setting
random_seed = args.randomseed
np.random.seed(random_seed)
torch.manual_seed(random_seed)
if torch.cuda.device_count() > 1:
torch.cuda.manual_seed_all(random_seed)
else:
torch.cuda.manual_seed(random_seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
self.slomo = model.Slomo(self.args.data_h, self.args.data_w, self.device)
self.slomo.to(self.device)
if self.args.init_type != "":
init_net(self.slomo, self.args.init_type)
print(self.args.init_type + " initializing slomo done!")
if self.args.train_continue:
if not self.args.nocomet and self.args.cometid != "":
self.comet_exp = ExistingExperiment(
previous_experiment=self.args.cometid
)
elif not self.args.nocomet and self.args.cometid == "":
self.comet_exp = Experiment(
workspace=self.args.workspace, project_name=self.args.projectname
)
else:
self.comet_exp = None
self.ckpt_dict = torch.load(self.args.checkpoint)
self.slomo.load_state_dict(self.ckpt_dict["model_state_dict"])
self.args.init_learning_rate = self.ckpt_dict["learningRate"]
self.optimizer = optim.Adam(
self.slomo.parameters(), lr=self.args.init_learning_rate
)
self.optimizer.load_state_dict(self.ckpt_dict["opt_state_dict"])
print("Pretrained model loaded!")
else:
# start logging info in comet-ml
if not self.args.nocomet:
self.comet_exp = Experiment(
workspace=self.args.workspace, project_name=self.args.projectname
)
# self.comet_exp.log_parameters(flatten_opts(self.args))
else:
self.comet_exp = None
self.ckpt_dict = {
"trainLoss": {},
"valLoss": {},
"valPSNR": {},
"valSSIM": {},
"learningRate": {},
"epoch": -1,
"detail": "End to end Super SloMo.",
"trainBatchSz": self.args.train_batch_size,
"validationBatchSz": self.args.validation_batch_size,
}
self.optimizer = optim.Adam(
self.slomo.parameters(), lr=self.args.init_learning_rate
)
self.scheduler = optim.lr_scheduler.MultiStepLR(
self.optimizer, milestones=self.args.milestones, gamma=0.1
)
# Channel wise mean calculated on adobe240-fps training dataset
mean = [0.5, 0.5, 0.5]
std = [1, 1, 1]
self.normalize = transforms.Normalize(mean=mean, std=std)
self.transform = transforms.Compose([transforms.ToTensor(), self.normalize])
trainset = dataloader.SuperSloMo(
root=self.args.dataset_root + "/train", transform=self.transform, train=True
)
self.trainloader = torch.utils.data.DataLoader(
trainset,
batch_size=self.args.train_batch_size,
num_workers=self.args.num_workers,
shuffle=True,
)
validationset = dataloader.SuperSloMo(
root=self.args.dataset_root + "/validation",
transform=self.transform,
# randomCropSize=(128, 128),
train=False,
)
self.validationloader = torch.utils.data.DataLoader(
validationset,
batch_size=self.args.validation_batch_size,
num_workers=self.args.num_workers,
shuffle=False,
)
### loss
self.supervisedloss = supervisedLoss()
self.best = {
"valLoss": 99999999,
"valPSNR": -1,
"valSSIM": -1,
}
self.checkpoint_counter = int(
(self.ckpt_dict["epoch"] + 1) / self.args.checkpoint_epoch
)
##[TensorboardX](https://github.com/lanpa/tensorboardX)
### For visualizing loss and interpolated frames
# writer = SummaryWriter("log")
###Initialize flow computation and arbitrary-time flow interpolation CNNs.
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
flowComp = model.UNet(6, 4)
flowComp.to(device)
if args.init_type != "":
init_net(flowComp, args.init_type)
print(args.init_type + " initializing flowComp done")
ArbTimeFlowIntrp = model.UNet(20, 5)
ArbTimeFlowIntrp.to(device)
if args.init_type != "":
init_net(ArbTimeFlowIntrp, args.init_type)
print(args.init_type + " initializing ArbTimeFlowIntrp done")
### Initialization
if args.train_continue:
if not args.nocomet and args.cometid != "":
comet_exp = ExistingExperiment(previous_experiment=args.cometid)
elif not args.nocomet and args.cometid == "":
image_file_names=val_file_names[::20],
to_augment=True,
transform=valid_transform)
train_loader = DataLoader(dataset=train_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=num_workers)
val_loader = DataLoader(dataset=val_dataset,
batch_size=batch_size,
shuffle=False,
num_workers=num_workers)
# Building Generator
netG = UNet_Colorization(num_classes=11, filters_base=16)
utils.init_net(netG)
summary(netG, input_size=(3, img_height, img_width))
# Building Discriminator
netD = Discriminator(input_nc=3,
img_height=img_height,
img_width=img_width,
filter_base=8,
num_block=5)
utils.init_net(netD)
summary(netD, input_size=(3, img_height, img_width))
# Optimizer
# G_optimizer = Adam(filter(lambda p: p.requires_grad, netG.parameters()), lr=lr, betas=(0.5, 0.999))
G_optimizer = Adam(netG.parameters(), lr=lr, betas=(0.5, 0.999))
D_optimizer = Adam(netD.parameters(), lr=lr, betas=(0.5, 0.999))
img_height=1024,
factor=0.01)
train_loader = DataLoader(dataset=train_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=num_workers)
lr = 1.0e-4
n_epochs = 200
add_log = False
add_output = True
model = UNet_softmax(num_classes=6,
filters_base=6,
input_channels=3,
add_output=add_output)
utils.init_net(model)
summary(model, input_size=(3, 480, 640))
optimizer = Adam(model.parameters(), lr=lr)
try:
model_root = root / "models"
model_root.mkdir(mode=0o777, parents=False)
except OSError:
print("")
try:
results_root = root / "results"
results_root.mkdir(mode=0o777, parents=False)
except OSError:
print("")
x, y, e = utils.get_data(f"data/{dataset}/{dataset}.npz", num_his, num_pred)
x, y, e = x.to(device), y.to(device), e.to(device)
n = x.shape[2]
adj = utils.get_adj("data/{}/distance.csv".format(dataset), n).to(device)
# adj = None
if not os.path.exists(f'experiment/{net_name}'):
os.mkdir(f'experiment/{net_name}')
net = AG_JNet(in_feature, dim_exp, layers_sm, layers_tm, n, num_his, num_pred,
device).to(device)
if net_name == "STGCN":
net = STGCN(n, in_feature, num_his, num_pred).to(device)
num_params = sum(param.numel() for param in net.parameters())
print('模型参数量：', num_params)
utils.init_net(net)
criterion = nn.MSELoss().to(device)
opt = Adam(net.parameters(), lr=lr)
num_sample = x.shape[0]
split_index = int(train_split * num_sample)
split_index1 = int(val_split * num_sample)
train_x, train_y, val_x, val_y, test_x, test_y, train_e, val_e, test_e, mean, std \
= utils.normalization(x, y, e, split_index, split_index1)
# 必须打乱数据集，不然loss降不下来。
train_loader = DataLoader(dataset=TensorDataset(train_x, train_e, train_y),
batch_size=batch_size,
shuffle=True)
val_loader = DataLoader(dataset=TensorDataset(val_x, val_e, val_y),
precompute_root.mkdir(mode=0o777, parents=True)
except OSError:
pass
feature_descriptor_model = models.FCDenseNet(
in_channels=3,
down_blocks=(3, 3, 3, 3, 3),
up_blocks=(3, 3, 3, 3, 3),
bottleneck_layers=4,
growth_rate=filter_growth_rate,
out_chans_first_conv=16,
feature_length=feature_length)
# Initialize the network with Kaiming He initialization
utils.init_net(feature_descriptor_model,
type="kaiming",
mode="fan_in",
activation_mode="relu",
distribution="normal")
# Multi-GPU running
feature_descriptor_model = torch.nn.DataParallel(feature_descriptor_model)
# Custom layer
response_map_generator = models.FeatureResponseGeneratorNoSoftThresholding(
)
# Evaluation metric
matching_accuracy_metric = losses.MatchingAccuracyMetric(threshold=3)
if trained_model_path.exists():
print("Loading {:s} ...".format(str(trained_model_path)))
pre_trained_state = torch.load(str(trained_model_path))
step = pre_trained_state['step']
elif opt.image_width == 128:
import models.vgg_128 as model
else:
raise ValueError('Unknown model: %s' % opt.model)
# define
frame_predictor = lstm_models.lstm((opt.factor+1)*opt.z_dim, opt.g_dim, opt.rnn_size, opt.predictor_rnn_layers, int(opt.batch_size/len(opt.gpu_ids)))
posterior_pose = lstm_models.gaussian_lstm(opt.g_dim+opt.factor*opt.z_dim, opt.z_dim, opt.rnn_size, opt.posterior_rnn_layers, int(opt.batch_size/len(opt.gpu_ids)))
prior = lstm_models.gaussian_lstm(opt.g_dim+opt.factor*opt.z_dim, opt.z_dim, opt.rnn_size, opt.prior_rnn_layers, int(opt.batch_size/len(opt.gpu_ids)))
cont_encoder = model.cont_encoder(opt.z_dim*opt.factor, opt.channels*opt.n_past) #g_dim = 64 or 128
pose_encoder = model.pose_encoder(opt.g_dim, opt.channels)
decoder = model.decoder(opt.g_dim, opt.channels)
# init
frame_predictor = utils.init_net(frame_predictor, init_type='normal', init_gain=0.02, gpu_ids=opt.gpu_ids)
posterior_pose = utils.init_net(posterior_pose, init_type='normal', init_gain=0.02, gpu_ids=opt.gpu_ids)
prior = utils.init_net(prior, init_type='normal', init_gain=0.02, gpu_ids=opt.gpu_ids)
cont_encoder = utils.init_net(cont_encoder, init_type='normal', init_gain=0.02, gpu_ids=opt.gpu_ids)
pose_encoder = utils.init_net(pose_encoder, init_type='normal', init_gain=0.02, gpu_ids=opt.gpu_ids)
decoder = utils.init_net(decoder, init_type='normal', init_gain=0.02, gpu_ids=opt.gpu_ids)
# load
utils.load_network(frame_predictor, 'frame_predictor', 'last', opt.model_path,device)
utils.load_network(posterior_pose, 'posterior_pose', 'last', opt.model_path,device)
utils.load_network(prior, 'prior', 'last', opt.model_path,device)
utils.load_network(cont_encoder, 'cont_encoder', 'last', opt.model_path,device)
utils.load_network(pose_encoder, 'pose_encoder', 'last', opt.model_path,device)
utils.load_network(decoder, 'decoder', 'last', opt.model_path,device)
# Fix randomness for reproducibility
torch.backends.cudnn.deterministic = False
torch.backends.cudnn.benchmark = True
torch.manual_seed(10085)
np.random.seed(10085)
random.seed(10085)
if not Path(args.precompute_root).exists():
Path(args.precompute_root).mkdir(parents=True)
model = models.FCDenseNetStd(
in_channels=3, down_blocks=(4, 4, 4, 4, 4),
up_blocks=(4, 4, 4, 4, 4), bottleneck_layers=4,
growth_rate=12, out_chans_first_conv=48)
# Initialize the depth estimation network with Kaiming He initialization
utils.init_net(model, type="kaiming", mode="fan_in", activation_mode="relu",
distribution="normal")
# Multi-GPU running
model = torch.nn.DataParallel(model)
# Load previous depth estimation model
if Path(args.trained_model_path).exists():
print("Loading {:s} ...".format(str(args.trained_model_path)))
state = torch.load(str(args.trained_model_path), encoding='latin1')
model.load_state_dict(state["model"])
step = state['step']
epoch = state['epoch']
print('Restored model, epoch {}, step {}'.format(epoch, step))
else:
print("No previous model detected")
raise OSError
label_path = "../datasets/lumi/B/test"
input_file_names = utils.read_lumi_filenames(input_path)
label_file_names = utils.read_lumi_filenames(label_path)
dataset = dataset.LumiDataset(input_filenames=input_file_names,
label_filenames=label_file_names,
transform=test_transform)
loader = DataLoader(dataset=dataset,
batch_size=batch_size,
shuffle=False,
num_workers=num_workers)
# Building Generator
netG = Generator(num_classes=3, filters_base=16)
netG = utils.init_net(netG)
summary(netG, input_size=(3, img_height, img_width))
try:
model_root = root / "models"
model_root.mkdir(mode=0o777, parents=False)
except OSError:
print("path exists")
try:
results_root = root / "results"
results_root.mkdir(mode=0o777, parents=False)
except OSError:
print("path exists")
print('Restoring mode')
def __init__(self, args=args):
super().__init__()
self.args = args
# random_seed setting
random_seed = args.randomseed
np.random.seed(random_seed)
torch.manual_seed(random_seed)
if torch.cuda.device_count() > 1:
torch.cuda.manual_seed_all(random_seed)
else:
torch.cuda.manual_seed(random_seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
self.pretrain_stage = self.args.pretrainstage
self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
self.slomofc = model.Slomofc(
self.args.data_h, self.args.data_w, self.device, self.pretrain_stage
)
self.slomofc.to(self.device)
if self.pretrain_stage:
self.learner = ContrastiveLearner(
self.slomofc,
image_size=128,
hidden_layer="avgpool",
use_momentum=True, # use momentum for key encoder
momentum_value=0.999,
project_hidden=False, # no projection heads
use_bilinear=True, # in paper, logits is bilinear product of query / key
use_nt_xent_loss=False, # use regular contrastive loss
augment_both=False, # in curl, only the key is augmented
)
if self.args.init_type != "":
init_net(self.slomofc, self.args.init_type)
print(self.args.init_type + " initializing slomo done!")
if self.args.train_continue:
if not self.args.nocomet and self.args.cometid != "":
self.comet_exp = ExistingExperiment(
previous_experiment=self.args.cometid
)
elif not self.args.nocomet and self.args.cometid == "":
self.comet_exp = Experiment(
workspace=self.args.workspace, project_name=self.args.projectname
)
else:
self.comet_exp = None
self.ckpt_dict = torch.load(self.args.checkpoint)
self.slomofc.load_state_dict(self.ckpt_dict["model_state_dict"])
self.args.init_learning_rate = self.ckpt_dict["learningRate"]
if not self.pretrain_stage:
self.optimizer = optim.Adam(
self.slomofc.parameters(), lr=self.args.init_learning_rate
)
else:
self.optimizer = optim.Adam(self.learner.parameters(), lr=3e-4)
self.optimizer.load_state_dict(self.ckpt_dict["opt_state_dict"])
print("Pretrained model loaded!")
else:
# start logging info in comet-ml
if not self.args.nocomet:
self.comet_exp = Experiment(
workspace=self.args.workspace, project_name=self.args.projectname
)
# self.comet_exp.log_parameters(flatten_opts(self.args))
else:
self.comet_exp = None
if not self.pretrain_stage:
self.ckpt_dict = {
"trainLoss": {},
"valLoss": {},
"valPSNR": {},
"valSSIM": {},
"learningRate": {},
"epoch": -1,
"detail": "End to end Super SloMo.",
"trainBatchSz": self.args.train_batch_size,
"validationBatchSz": self.args.validation_batch_size,
}
else:
self.ckpt_dict = {
"conLoss": {},
"learningRate": {},
"epoch": -1,
"detail": "Pretrain_stage of Super SloMo.",
"trainBatchSz": self.args.train_batch_size,
}
if not self.pretrain_stage:
self.optimizer = optim.Adam(
self.slomofc.parameters(), lr=self.args.init_learning_rate
)
else:
self.optimizer = optim.Adam(self.learner.parameters(), lr=3e-4)
self.scheduler = optim.lr_scheduler.MultiStepLR(
self.optimizer, milestones=self.args.milestones, gamma=0.1
)
# Channel wise mean calculated on adobe240-fps training dataset
if not self.pretrain_stage:
mean = [0.5, 0.5, 0.5]
std = [1, 1, 1]
self.normalize = transforms.Normalize(mean=mean, std=std)
self.transform = transforms.Compose([transforms.ToTensor(), self.normalize])
else:
self.transform = transforms.Compose([transforms.ToTensor()])
trainset = dataloader.SuperSloMo(
root=self.args.dataset_root + "/train", transform=self.transform, train=True
)
self.trainloader = torch.utils.data.DataLoader(
trainset,
batch_size=self.args.train_batch_size,
num_workers=self.args.num_workers,
shuffle=True,
)
if not self.pretrain_stage:
validationset = dataloader.SuperSloMo(
root=self.args.dataset_root + "/validation",
transform=self.transform,
# randomCropSize=(128, 128),
train=False,
)
self.validationloader = torch.utils.data.DataLoader(
validationset,
batch_size=self.args.validation_batch_size,
num_workers=self.args.num_workers,
shuffle=False,
)
### loss
if not self.pretrain_stage:
self.supervisedloss = supervisedLoss()
self.best = {
"valLoss": 99999999,
"valPSNR": -1,
"valSSIM": -1,
}
else:
self.best = {
"conLoss": 99999999,
}
self.checkpoint_counter = int(
(self.ckpt_dict["epoch"] + 1) / self.args.checkpoint_epoch
)
