Resize((256, 128)),
ToTensor(),
Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
print('Start dataloader...')
train_dataloader = utils.Get_Video_train_DataLoader(args.train_txt,args.train_info, train_transform, shuffle=True,num_workers=args.num_workers,\
S=args.S,track_per_class=args.track_per_class,class_per_batch=args.class_per_batch)
num_class = train_dataloader.dataset.n_id
test_dataloader = utils.Get_Video_test_DataLoader(args.test_txt,args.test_info,args.query_info,test_transform,batch_size=args.batch_size,\
shuffle=False,num_workers=args.num_workers,S=args.S,distractor=True)
print('End dataloader...')
network = nn.DataParallel(
models.CNN(args.latent_dim,
model_type=args.model_type,
num_class=num_class,
non_layers=args.non_layers,
stripes=args.stripes,
temporal=args.temporal).cuda())
if args.load_ckpt is not None:
state = torch.load(args.load_ckpt)
network.load_state_dict(state, strict=False)
# log
os.system('mkdir -p %s' % (args.ckpt))
f = open(os.path.join(args.ckpt, args.log_path), 'a')
f.close()
# Train loop
# 1. Criterion
criterion_triplet = TripletLoss('soft', True)
critetion_id = nn.CrossEntropyLoss().cuda()
Resize((256, 128)),
ToTensor(),
Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
print('Start dataloader...')
train_dataloader = utils.Get_Video_train_DataLoader(args.train_txt,args.train_info, train_transform, shuffle=True,num_workers=args.num_workers,\
S=args.S,track_per_class=args.track_per_class,class_per_batch=args.class_per_batch)
num_class = train_dataloader.dataset.n_id
test_dataloader = utils.Get_Video_test_DataLoader(args.test_txt,args.test_info,args.query_info,test_transform,batch_size=args.batch_size,\
shuffle=False,num_workers=args.num_workers,S=args.S,distractor=True)
print('End dataloader...\n')
network = nn.DataParallel(
models.CNN(args.latent_dim,
model_type=args.model_type,
num_class=num_class,
stride=args.stride).cuda())
if args.load_ckpt is not None:
state = torch.load(args.load_ckpt)
network.load_state_dict(state)
# log
os.system('mkdir -p %s' % (args.ckpt))
f = open(os.path.join(args.ckpt, args.log_path), 'a')
f.close()
# Train loop
# 1. Criterion
criterion_triplet = TripletLoss('soft', True)
criterion_ID = nn.CrossEntropyLoss().cuda()