current_device = torch.cuda.current_device()
print("Running on", torch.cuda.get_device_name(current_device))
else:
print("Running on CPU")
data_transforms = get_transforms(
input_size=(args.size, args.size),
image_mode=False
)
train_dataset = get_datasets(
name_list=["mydataset"],
split_list=["train"],
config_path=args.dataset_config,
root=args.data,
training=True,
transforms=data_transforms['train'],
read_clip=True,
random_reverse_clip=True,
clip_len=args.clip_len
)
val_dataset = get_datasets(
name_list="mydataset",
split_list="val",
config_path=args.dataset_config,
root=args.data,
training=True,
transforms=data_transforms['val'],
read_clip=True,
random_reverse_clip=False,
clip_len=args.clip_len
input_size=(args.size, args.size),
image_mode=True
)
# data_transformsWO = get_transformsWO(
# input_size=(args.size, args.size),
# image_mode=False
# )
train_dataset = get_datasets(
# name_list=["DAVIS2016", "FBMS", "VOS"],
# name_list=["DAVIS2016"], #先只用DAVIS2016-test做测试
name_list=["DAVIS2016", "DAVSOD", "VOS"], # 上下得一起改
split_list=["train", "train", "train"],
# name_list=["FBMS"],
# split_list=["train"],
config_path=args.dataset_config,
root=args.data,
training=True,
transforms=data_transforms['train'],
read_clip=True,
random_reverse_clip=False, # 是否反转clip
clip_len=args.clip_len
)
val_dataset = get_datasets(
name_list=["DAVIS2016"],
split_list=["val"],
config_path=args.dataset_config,
root=args.data,
training=True,
transforms=data_transforms['val'],
current_device = torch.cuda.current_device()
print("Running on", torch.cuda.get_device_name(current_device))
else:
print("Running on CPU")
data_transforms = get_transforms(
input_size=(args.size, args.size),
image_mode=False
)
dataset = get_datasets(
name_list=args.dataset,
split_list=args.split,
config_path=args.dataset_config,
root=args.data,
training=False,
transforms=data_transforms['test'],
read_clip=True,
random_reverse_clip=False,
label_interval=1,
frame_between_label_num=0,
clip_len=args.clip_len
)
dataloader = data.DataLoader(
dataset=dataset,
batch_size=1, # only support 1 video clip
num_workers=args.num_workers,
shuffle=False
)
model = VideoModel(output_stride=args.os)
if cuda:
torch.backends.cudnn.benchmark = True
current_device = torch.cuda.current_device()
print("Running on", torch.cuda.get_device_name(current_device))
else:
print("Running on CPU")
data_transforms = get_transforms(input_size=(args.size, args.size),
image_mode=False)
train_dataset = get_datasets(name_list=["DAVIS2016", "FBMS", "VOS"],
split_list=["train", "train", "train"],
config_path=args.dataset_config,
root=args.data,
training=True,
transforms=data_transforms['train'],
read_clip=True,
random_reverse_clip=True,
label_interval=args.label_interval,
clip_len=args.clip_len)
val_dataset = get_datasets(name_list="VOS",
split_list="val",
config_path=args.dataset_config,
root=args.data,
training=True,
transforms=data_transforms['val'],
read_clip=True,
random_reverse_clip=False,
label_interval=args.label_interval,
clip_len=args.clip_len)