def train_model(dataset=dataset, save_dir=save_dir, num_classes=num_classes, lr=lr, num_epochs=nEpochs, save_epoch=snapshot, useTest=useTest, test_interval=nTestInterval): """ Args: num_classes (int): Number of classes in the data num_epochs (int, optional): Number of epochs to train for. """ if modelName == 'C3D': model = C3D_model.C3D(num_classes=num_classes, pretrained=True) train_params = [{ 'params': C3D_model.get_1x_lr_params(model), 'lr': lr }, { 'params': C3D_model.get_10x_lr_params(model), 'lr': lr * 10 }] elif modelName == 'R2Plus1D': model = R2Plus1D_model.R2Plus1DClassifier(num_classes=num_classes, layer_sizes=(2, 2, 2, 2)) train_params = [{ 'params': R2Plus1D_model.get_1x_lr_params(model), 'lr': lr }, { 'params': R2Plus1D_model.get_10x_lr_params(model), 'lr': lr * 10 }] else: print('We only implemented C3D and R2Plus1D models.') raise NotImplementedError criterion = nn.CrossEntropyLoss( ) # standard crossentropy loss for classification optimizer = optim.SGD(train_params, lr=lr, momentum=0.9, weight_decay=5e-4) scheduler = optim.lr_scheduler.StepLR( optimizer, step_size=10, gamma=0.1) # the scheduler divides the lr by 10 every 10 epochs if resume_epoch == 0: print("Training {} from scratch...".format(modelName)) else: checkpoint = torch.load(os.path.join( save_dir, 'models', modelName + '_' + dataset + '_epoch-' + str(resume_epoch - 1) + '.pth.tar'), map_location=lambda storage, loc: storage ) # Load all tensors onto the CPU print("Initializing weights from: {}...".format( os.path.join( save_dir, 'models', modelName + '_' + dataset + '_epoch-' + str(resume_epoch - 1) + '.pth.tar'))) model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['opt_dict']) print('Total params: %.2fM' % (sum(p.numel() for p in model.parameters()) / 1000000.0)) model.to(device) criterion.to(device) log_dir = os.path.join( save_dir, 'models', datetime.now().strftime('%b%d_%H-%M-%S') + '_' + socket.gethostname()) writer = SummaryWriter(log_dir=log_dir) print('Training model on {} dataset...'.format(dataset)) train_dataloader = DataLoader(VideoDataset(dataset=dataset, split='train', clip_len=16), batch_size=16, shuffle=True, num_workers=4) val_dataloader = DataLoader(VideoDataset(dataset=dataset, split='val', clip_len=16), batch_size=16, num_workers=4) test_dataloader = DataLoader(VideoDataset(dataset=dataset, split='test', clip_len=16), batch_size=16, num_workers=4) trainval_loaders = {'train': train_dataloader, 'val': val_dataloader} trainval_sizes = { x: len(trainval_loaders[x].dataset) for x in ['train', 'val'] } test_size = len(test_dataloader.dataset) for epoch in range(resume_epoch, num_epochs): # each epoch has a training and validation step for phase in ['train', 'val']: start_time = timeit.default_timer() # reset the running loss and corrects running_loss = 0.0 running_corrects = 0 # set model to train() or eval() mode depending on whether it is trained # or being validated. Primarily affects layers such as BatchNorm or Dropout. if phase == 'train': # scheduler.step() is to be called once every epoch during training scheduler.step() model.train() else: model.eval() for inputs, labels in tqdm(trainval_loaders[phase]): # move inputs and labels to the device the training is taking place on inputs = Variable(inputs, requires_grad=True).to(device) labels = Variable(labels).to(device) optimizer.zero_grad() outputs = model(inputs) probs = nn.Softmax(dim=1)(outputs) preds = torch.max(probs, 1)[1] loss = criterion(outputs, labels) if phase == 'train': loss.backward() optimizer.step() running_loss += loss.item() * inputs.size(0) running_corrects += torch.sum(preds == labels.data) epoch_loss = running_loss / trainval_sizes[phase] epoch_acc = running_corrects.double() / trainval_sizes[phase] if phase == 'train': writer.add_scalar('data/train_loss_epoch', epoch_loss, epoch) writer.add_scalar('data/train_acc_epoch', epoch_acc, epoch) else: writer.add_scalar('data/val_loss_epoch', epoch_loss, epoch) writer.add_scalar('data/val_acc_epoch', epoch_acc, epoch) print("[{}] Epoch: {}/{} Loss: {} Acc: {}".format( phase, epoch + 1, nEpochs, epoch_loss, epoch_acc)) stop_time = timeit.default_timer() print("Execution time: " + str(stop_time - start_time) + "\n") if epoch % save_epoch == (save_epoch - 1): torch.save( { 'epoch': epoch + 1, 'state_dict': model.state_dict(), 'opt_dict': optimizer.state_dict(), }, os.path.join( save_dir, 'models', modelName + '_' + dataset + '_epoch-' + str(epoch) + '.pth.tar')) print("Save model at {}\n".format( os.path.join( save_dir, 'models', modelName + '_' + dataset + '_epoch-' + str(epoch) + '.pth.tar'))) if useTest and epoch % test_interval == (test_interval - 1): model.eval() start_time = timeit.default_timer() running_loss = 0.0 running_corrects = 0 for inputs, labels in tqdm(test_dataloader): inputs = inputs.to(device) labels = labels.to(device) with torch.no_grad(): outputs = model(inputs) probs = nn.Softmax(dim=1)(outputs) preds = torch.max(probs, 1)[1] loss = criterion(outputs, labels) running_loss += loss.item() * inputs.size(0) running_corrects += torch.sum(preds == labels.data) epoch_loss = running_loss / test_size epoch_acc = running_corrects.double() / test_size writer.add_scalar('data/test_loss_epoch', epoch_loss, epoch) writer.add_scalar('data/test_acc_epoch', epoch_acc, epoch) print("[test] Epoch: {}/{} Loss: {} Acc: {}".format( epoch + 1, nEpochs, epoch_loss, epoch_acc)) stop_time = timeit.default_timer() print("Execution time: " + str(stop_time - start_time) + "\n") writer.close()
def train_model(dataset=dataset, save_dir=SAVE_FILE_FOLDER, num_classes=num_classes, lr=lr, num_epochs=nEpochs, save_epoch=snapshot, useTest=useTest, test_interval=nTestInterval): """ Args: num_classes (int): Number of classes in the data num_epochs (int, optional): Number of epochs to train for. """ if modelName == 'C3D': model = C3D_model.C3D(num_classes=num_classes, pretrained=IF_PRETRAIN) train_params = [{'params': C3D_model.get_1x_lr_params(model), 'lr': lr}, {'params': C3D_model.get_10x_lr_params(model), 'lr': lr * 10}] elif modelName == 'C3D_td5': model = C3D_model.C3D_td5(num_classes=num_classes, pretrained=IF_PRETRAIN) train_params = [{'params': C3D_model.get_1x_lr_params(model), 'lr': lr}, {'params': C3D_model.get_10x_lr_params(model), 'lr': lr * 10}] elif modelName == 'R2Plus1D': model = R2Plus1D_model.R2Plus1DClassifier(num_classes=num_classes, layer_sizes=(2, 2, 2, 2)) train_params = [{'params': R2Plus1D_model.get_1x_lr_params(model), 'lr': lr}, {'params': R2Plus1D_model.get_10x_lr_params(model), 'lr': lr * 10}] elif modelName == 'R3D': model = R3D_model.R3DClassifier(num_classes=num_classes, layer_sizes=(2, 2, 2, 2)) train_params = model.parameters() else: print('We only implemented C3D and R2Plus1D models.') raise NotImplementedError criterion = nn.CrossEntropyLoss() # standard crossentropy loss for classification if _optimizer == "SGD": optimizer = optim.SGD(train_params, lr=lr, momentum=MOMENTUM, weight_decay=WD) elif _optimizer == "Adam": optimizer = optim.Adam(train_params, lr=lr, weight_decay=WD) # print(optimizer) scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=SCHEDULER_STEP_SIZE, gamma=SCHEDULER_GAMMA) # the scheduler divides the lr by 10 every 10 epochs model.to(device) criterion.to(device) # if resume_epoch == 0: if resume_model_path == None: print("Training {} from scratch...".format(modelName)) else: checkpoint = torch.load( resume_model_path, map_location=lambda storage, loc: storage) # Load all tensors onto the CPU print("Initializing weights from: {}...".format(resume_model_path)) model.load_state_dict(checkpoint['state_dict']) if RESUM_OPTIMIZER: optimizer.load_state_dict(checkpoint['opt_dict']) # resume_epoch # else: # checkpoint = torch.load(os.path.join(SAVE_FILE_FOLDER, 'models', EXP_NAME + '_epoch-' + str(resume_epoch - 1) + '.pth.tar'), # map_location=lambda storage, loc: storage) # Load all tensors onto the CPU # print("Initializing weights from: {}...".format( # os.path.join(SAVE_FILE_FOLDER, EXP_NAME + '_epoch-' + str(resume_epoch - 1) + '.pth.tar'))) # model.load_state_dict(checkpoint['state_dict']) # optimizer.load_state_dict(checkpoint['opt_dict']) print('Total params: %.2fM' % (sum(p.numel() for p in model.parameters()) / 1000000.0)) model.to(device) criterion.to(device) writer = SummaryWriter(logdir=LOG_PATH) print('Training model on {} dataset...'.format(dataset)) train_dataloader = DataLoader(VideoDataset(dataset=dataset, split='train', clip_len=clip_len, preprocess=IF_PREPROCESS_TRAIN, grayscale=grayscale), batch_size=BS, shuffle=True, num_workers=N_WORKERS) val_dataloader = DataLoader(VideoDataset(dataset=dataset, split='val', clip_len=clip_len, preprocess=IF_PREPROCESS_VAL, grayscale=grayscale), batch_size=BS, num_workers=N_WORKERS) test_dataloader = DataLoader(VideoDataset(dataset=dataset, split='test', clip_len=clip_len, preprocess=IF_PREPROCESS_TEST, grayscale=grayscale), batch_size=BS, num_workers=N_WORKERS) trainval_loaders = {'train': train_dataloader, 'val': val_dataloader} trainval_sizes = {x: len(trainval_loaders[x].dataset) for x in ['train', 'val']} test_size = len(test_dataloader.dataset) cudnn.benchmark = True global_best_val_acc = 0 for epoch in range(num_epochs): # each epoch has a training and validation step for phase in ['train', 'val']: start_time = timeit.default_timer() # reset the running loss and corrects running_loss = 0.0 running_corrects = 0.0 # running_roc = 0.0 list_pred = list() list_label = list() # print(optimizer) # set model to train() or eval() mode depending on whether it is trained # or being validated. Primarily affects layers such as BatchNorm or Dropout. if phase == 'train': # scheduler.step() is to be called once every epoch during training scheduler.step() model.train() else: model.eval() # for inputs, labels in tqdm(trainval_loaders[phase]): run_count = 0 for inputs, labels in trainval_loaders[phase]: # move inputs and labels to the device the training is taking place on inputs = Variable(inputs, requires_grad=True).to(device) labels = Variable(labels).to(device) optimizer.zero_grad() if phase == 'train': outputs = model(inputs) else: with torch.no_grad(): outputs = model(inputs) probs = nn.Softmax(dim=1)(outputs) preds = torch.max(probs, 1)[1] loss = criterion(outputs, labels) if phase == 'train': loss.backward() optimizer.step() running_loss += loss.item() * inputs.size(0) running_corrects += torch.sum(preds == labels.data) # try: # running_roc += roc_auc_score(labels.data.cpu(), preds.cpu()) # except: # y_true = labels.data.cpu().tolist() # y_true_2 = y_true.copy() # for i_cls in range(N_CLASSES): # y_true_2.append(i_cls) # # y_pred = preds.cpu().tolist() # y_pred_2 = y_pred.copy() # for i_cls in range(N_CLASSES): # y_pred_2.append(i_cls) # # running_roc += roc_auc_score(y_true_2, y_pred_2) # # run_count += 1 list_label += labels.data.cpu().tolist() list_pred += preds.cpu().tolist() epoch_loss = running_loss / trainval_sizes[phase] epoch_acc = running_corrects.double() / trainval_sizes[phase] epoch_roc = multiclass_roc_score(label=list_label, pred=list_pred, n_cls=N_CLASSES) if phase == 'train': writer.add_scalar('data/train_loss_epoch', epoch_loss, epoch) writer.add_scalar('data/train_acc_epoch', epoch_acc, epoch) writer.add_scalar('data/train_roc_epoch', epoch_roc, epoch) else: writer.add_scalar('data/val_loss_epoch', epoch_loss, epoch) writer.add_scalar('data/val_acc_epoch', epoch_acc, epoch) writer.add_scalar('data/val_roc_epoch', epoch_roc, epoch) # if epoch_acc >= global_best_val_acc: # torch.save({ # 'epoch': epoch + 1, # 'state_dict': model.state_dict(), # 'opt_dict': optimizer.state_dict(), # }, os.path.join(SAVE_FILE_FOLDER, 'models', EXP_NAME + '_epoch-' + str(epoch) + 'ValAcc_{:10.4f}_'.format(epoch_loss) + '.pth.tar')) # print("Save model at {}\n".format( # os.path.join(SAVE_FILE_FOLDER, 'models', EXP_NAME + '_epoch-' + str(epoch) + 'ValAcc_{:10.4f}_'.format(epoch_loss) + '.pth.tar'))) print("[{}] Epoch: {}/{} Loss: {} Acc: {}, ROC:{}".format(phase, epoch+1, nEpochs, epoch_loss, epoch_acc, epoch_roc)) stop_time = timeit.default_timer() # print("Execution time: " + str(stop_time - start_time) + "\n") if epoch % save_epoch == (save_epoch - 1): torch.save({ 'epoch': epoch + 1, 'state_dict': model.state_dict(), 'opt_dict': optimizer.state_dict(), }, os.path.join(SAVE_FILE_FOLDER, EXP_NAME + '_epoch-' + str(epoch) + '.pth.tar')) print("Save model at {}\n".format(os.path.join(SAVE_FILE_FOLDER, EXP_NAME + '_epoch-' + str(epoch) + '.pth.tar'))) if useTest and epoch % test_interval == (test_interval - 1): model.eval() start_time = timeit.default_timer() running_loss = 0.0 running_corrects = 0.0 # running_roc = 0.0 list_pred = list() list_label = list() # for inputs, labels in tqdm(test_dataloader): run_count = 0 for inputs, labels in test_dataloader: inputs = inputs.to(device) labels = labels.to(device) with torch.no_grad(): outputs = model(inputs) probs = nn.Softmax(dim=1)(outputs) preds = torch.max(probs, 1)[1] loss = criterion(outputs, labels) running_loss += loss.item() * inputs.size(0) running_corrects += torch.sum(preds == labels.data) # try: # running_roc += roc_auc_score(labels.data.cpu(), preds.cpu()) # except: # running_roc += 0.5 # run_count += 1 list_label += labels.data.cpu().tolist() list_pred += preds.cpu().tolist() epoch_loss = running_loss / test_size epoch_acc = running_corrects.double() / test_size epoch_roc = multiclass_roc_score(label=list_label, pred=list_pred, n_cls=N_CLASSES) writer.add_scalar('data/test_loss_epoch', epoch_loss, epoch) writer.add_scalar('data/test_acc_epoch', epoch_acc, epoch) writer.add_scalar('data/test_roc_epoch', epoch_roc, epoch) print("[test] Epoch: {}/{} Loss: {} Acc:{} ROC: {}".format(epoch+1, nEpochs, epoch_loss, epoch_acc, epoch_roc)) stop_time = timeit.default_timer() # print("Execution time: " + str(stop_time - start_time) + "\n") writer.close()
def train_model(dataset=dataset, save_dir=save_dir, num_classes=num_classes, lr=lr, num_epochs=nEpochs, save_epoch=snapshot, useTest=useTest, test_interval=nTestInterval): """ Args: num_classes (int): Number of classes in the data num_epochs (int, optional): Number of epochs to train for. """ if modelName == 'C3D': model = C3D_model.C3D(num_classes=num_classes, pretrained=False) train_params = [{'params': C3D_model.get_1x_lr_params(model), 'lr': lr}, {'params': C3D_model.get_10x_lr_params(model), 'lr': lr * 10}] elif modelName == 'R2Plus1D': model = R2Plus1D_model.R2Plus1DClassifier(num_classes=num_classes, layer_sizes=(2, 2, 2, 2)) train_params = [{'params': R2Plus1D_model.get_1x_lr_params(model), 'lr': lr}, {'params': R2Plus1D_model.get_10x_lr_params(model), 'lr': lr * 10}] elif modelName == 'R3D': model = R3D_model.R3DClassifier(num_classes=num_classes, layer_sizes=(3, 4, 6, 3)) # model = resnet.ResNet(num_classes=num_classes, layers=(3, 4, 6, 3), sample_size=112, sample_duration=16) train_params = model.parameters() elif modelName == 'R2D': model = R2Dnet.R2DClassifier(group_num_classes=num_classes, pretrained=True) # model = resnet.ResNet(num_classes=num_classes, layers=(3, 4, 6, 3), sample_size=112, sample_duration=16) train_params = model.parameters() else: print('We only implemented C3D and R2Plus1D models.') raise NotImplementedError criterion = nn.CrossEntropyLoss() # standard crossentropy loss for classification optimizer = optim.SGD(train_params, lr=lr, momentum=0.9, weight_decay=5e-4) scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1) # the scheduler divides the lr by 10 every 10 epochs model.to(device) #move here because resume need .cuda() if resume_epoch == 0: print("Training {} from scratch...".format(modelName)) else: checkpoint = torch.load(os.path.join(save_dir, 'models', saveName + '_epoch-' + str(resume_epoch - 1) + '.pth.tar'), map_location=lambda storage, loc: storage) # Load all tensors onto the CPU print("Initializing weights from: {}...".format( os.path.join(save_dir, 'models', saveName + '_epoch-' + str(resume_epoch - 1) + '.pth.tar'))) model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['opt_dict']) print('Total params: %.2fM' % (sum(p.numel() for p in model.parameters()) / 1000000.0)) # model.to(device) criterion.to(device) log_dir = os.path.join(save_dir, 'models', datetime.now().strftime('%b%d_%H-%M-%S') + '_' + socket.gethostname()) writer = SummaryWriter(log_dir=log_dir) print('Training model on {} dataset...'.format(dataset)) train_dataloader = DataLoader(VolleyballDataset(dataset=dataset, split='train',clip_len=16), batch_size=4, shuffle=True, \ num_workers=0) val_dataloader = DataLoader(VolleyballDataset(dataset=dataset, split='val', clip_len=16), batch_size=4, num_workers=0) test_dataloader = DataLoader(VolleyballDataset(dataset=dataset, split='test', clip_len=16), batch_size=4, num_workers=0) trainval_loaders = {'train': train_dataloader, 'val': val_dataloader} trainval_sizes = {x: len(trainval_loaders[x].dataset) for x in ['train', 'val']} test_size = len(test_dataloader.dataset) for epoch in range(resume_epoch, num_epochs): # each epoch has a training and validation step for phase in ['train', 'val']: start_time = timeit.default_timer() # reset the running loss and corrects running_loss = 0.0 running_corrects = 0.0 # set model to train() or eval() mode depending on whether it is trained # or being validated. Primarily affects layers such as BatchNorm or Dropout. if phase == 'train': # scheduler.step() is to be called once every epoch during training scheduler.step() model.train() else: model.eval() torch.backends.cudnn.benchmark=False # for inputs, bbox_inputs, labels, adjacent_matrix in tqdm(trainval_loaders[phase]): # for inputs, labels in tqdm(trainval_loaders[phase]): for inputs, labels, dists in tqdm(trainval_loaders[phase]): # move inputs and labels to the device the training is taking place on inputs = Variable(inputs, requires_grad=True).to(device) # bbox_inputs = Variable(bbox_inputs, requires_grad=True).to(device) # adjacent_matrix = Variable(adjacent_matrix, requires_grad=True).to(device) labels = Variable(labels).to(device) dists = Variable(dists, requires_grad = True).to(device) # dist_num = Variable(dist_num).to(device) optimizer.zero_grad() if phase == 'train': outputs = model(inputs, dists) # outputs = model(inputs, bbox_inputs, adjacent_matrix) else: with torch.no_grad(): # outputs = model(inputs, bbox_inputs, adjacent_matrix) outputs = model(inputs, dists) probs = nn.Softmax(dim=1)(outputs) preds = torch.max(probs, 1)[1] loss = criterion(outputs, labels) print("labels",labels) print("outputs",outputs) print("loss",loss) torch.backends.cudnn.benchmark = False if phase == 'train': loss.backward(retain_graph=True) optimizer.step() running_loss += loss.item() * inputs.size(0) running_corrects += torch.sum(preds == labels.data) epoch_loss = running_loss / trainval_sizes[phase] epoch_acc = running_corrects.double() / trainval_sizes[phase] if phase == 'train': writer.add_scalar('data/train_loss_epoch', epoch_loss, epoch) writer.add_scalar('data/train_acc_epoch', epoch_acc, epoch) else: writer.add_scalar('data/val_loss_epoch', epoch_loss, epoch) writer.add_scalar('data/val_acc_epoch', epoch_acc, epoch) print("[{}] Epoch: {}/{} Loss: {} Acc: {}".format(phase, epoch+1, nEpochs, epoch_loss, epoch_acc)) stop_time = timeit.default_timer() print("Execution time: " + str(stop_time - start_time) + "\n") if epoch % save_epoch == (save_epoch - 1): print(os.path.join(save_dir, 'models', saveName + \ '_epoch-' + str(epoch) + '.pth.tar')) torch.save({ 'epoch': epoch + 1, 'state_dict': model.state_dict(), 'opt_dict': optimizer.state_dict(), }, os.path.join(save_dir, 'models', saveName + '_epoch-' + str(epoch) + '.pth.tar')) print("Save model at {}\n".format(os.path.join(save_dir, 'models', saveName + '_epoch-' + str(epoch) + '.pth.tar'))) if useTest and epoch % test_interval == (test_interval - 1): model.eval() start_time = timeit.default_timer() running_loss = 0.0 running_corrects = 0.0 for inputs, bbox_inputs, labels, adjacent_matrix in tqdm(test_dataloader): # for inputs, labels in tqdm(test_dataloader): bbox_inputs = Variable(bbox_inputs, requires_grad=True).to(device) adjacent_matrix = Variable(adjacent_matrix, requires_grad=True).to(device) inputs = Variable(inputs.to(device)) with torch.no_grad(): # outputs = model(inputs) outputs = model(inputs, bbox_inputs, adjacent_matrix) probs = nn.Softmax(dim=1)(outputs) preds = torch.max(probs, 1)[1] loss = criterion(outputs, labels) running_loss += loss.item() * inputs.size(0) running_corrects += torch.sum(preds == labels.data) epoch_loss = running_loss / test_size epoch_acc = running_corrects.double() / test_size writer.add_scalar('data/test_loss_epoch', epoch_loss, epoch) writer.add_scalar('data/test_acc_epoch', epoch_acc, epoch) print("[test] Epoch: {}/{} Loss: {} Acc: {}".format(epoch+1, nEpochs, epoch_loss, epoch_acc)) stop_time = timeit.default_timer() print("Execution time: " + str(stop_time - start_time) + "\n") writer.close()
def test_model(dataset=dataset, save_dir=save_dir, num_classes=num_classes, lr=lr, save_epoch=snapshot, useTest=useTest, test_interval=nTestInterval): """ Args: num_classes (int): Number of classes in the data num_epochs (int, optional): Number of epochs to train for. """ if modelName == 'C3DVA': model = C3DVA_model.C3DVA(num_classes=2, pretrained=True) train_params = model.parameters() elif modelName == 'CSN': model = CSN_model.csn26(num_classes=2, add_landmarks=ADD_LANDMARKS) train_params = model.parameters() elif modelName == 'R2Plus1D': model = R2Plus1D_model.R2Plus1DRegressor(num_classes=2, layer_sizes=LAYER_SIZES) train_params = [{'params': R2Plus1D_model.get_1x_lr_params(model), 'lr': lr}, {'params': R2Plus1D_model.get_10x_lr_params(model), 'lr': lr * 10}] elif modelName == 'Transformer': model = transformer_v3.Semi_Transformer(num_classes=40, seq_len=CLIP_LEN) train_params = model.parameters() elif modelName == 'Resnet3d': model = resnet3D.resnet3d18(num_classes=400, pretrained=None) train_params = model.parameters() checkpoint = torch.load(os.path.join(save_dir, 'models', saveName + '_epoch-' + str(resume_epoch - 1) + '.pth.tar'), map_location=lambda storage, loc: storage) # Load all tensors onto the CPU print("Initializing weights from: {}...".format( os.path.join(save_dir, 'models', saveName + '_epoch-' + str(resume_epoch - 1) + '.pth.tar'))) model.load_state_dict(checkpoint['state_dict'], strict=False) print('Total params: %.2fM' % (sum(p.numel() for p in model.parameters()) / 1000000.0)) model.to(device) print('Training model on {} dataset...'.format(dataset)) test_dataloader = DataLoader(VideoDataset(dataset=dataset, label_type="VA_Set", split=split, \ clip_len=CLIP_LEN, stride=FRAME_STRIDE, add_landmarks=ADD_LANDMARKS, triplet_label=True), batch_size=1, shuffle=False, num_workers=0,\ drop_last=True, pin_memory=False) test_size = len(test_dataloader.dataset) # testing model.eval() start_time = timeit.default_timer() criterion = SetVACriterion(num_classes=20, use_mse=False, is_test=True) criterion.to(device) val_cccs = [] aro_cccs = [] val_mses = [] aro_mses = [] with open('test_log.txt', 'a') as test_log: test_log.write('=======start=======\n') test_log.write('model path: ' + os.path.join(save_dir, 'models', saveName + '_epoch-' + str(resume_epoch - 1) + '.pth.tar') + '\n') for frames, label_path in (test_dataloader): print(frames[0][0].split('/')[-2]) frame_ids = [] # 图片对应的帧号 for each in frames: fid = int(each[0].split('/')[-1].split('.')[0]) frame_ids.append(fid) frame_count = len(frames) frame_root_dir = '/' + os.path.join(*frames[0][0].split('/')[1:-1]) # valence_label = np.empty((frame_count-CLIP_LEN*FRAME_STRIDE+1, 1), np.dtype('float32')) # arousal_label = np.empty((frame_count-CLIP_LEN*FRAME_STRIDE+1, 1), np.dtype('float32')) # preds_all = np.empty((frame_count-CLIP_LEN*FRAME_STRIDE+1, 2), np.dtype('float32')) valence_label = np.zeros((frame_ids[-1], 1), np.dtype('float32')) arousal_label = np.zeros((frame_ids[-1], 1), np.dtype('float32')) preds_all = np.zeros((frame_ids[-1], 2), np.dtype('float32')) '''读取视频的所有帧''' for center in tqdm( frame_ids ): if ADD_LANDMARKS: buffer = np.empty(( CLIP_LEN, RESIZE_HEIGHT, RESIZE_WIDTH, 4), np.dtype('float32')) else: buffer = np.empty(( CLIP_LEN, RESIZE_HEIGHT, RESIZE_WIDTH, 3), np.dtype('float32')) '''读取图像''' clip = list(range(center-int(CLIP_LEN/2)*FRAME_STRIDE, center+int(CLIP_LEN/2)*FRAME_STRIDE, FRAME_STRIDE)) for i, frame_id in enumerate(clip): frame_name = None if (frame_id < 1 or frame_id > frame_ids[-1]): frame = np.ones((RESIZE_HEIGHT, RESIZE_WIDTH, 3))*128 else: frame_name = os.path.join(frame_root_dir,str(frame_id).zfill(5)+'.jpg') #print(center, frames[i]) if os.path.exists(frame_name): frame = np.array(cv2.imread(frame_name)).astype(np.float32) else: frame = np.ones((RESIZE_HEIGHT, RESIZE_WIDTH, 3))*128 frame = cv2.resize(frame, (RESIZE_HEIGHT, RESIZE_WIDTH)) # cv2.imshow(frame_name, frame/256) # cv2.waitKey(1000) # cv2.destroyAllWindows() if ADD_LANDMARKS: # 读landmarks if frame_name != None: lm_name = frame_name.replace('image', 'landmarks') lm_img = np.array(cv2.imread(lm_name, cv2.IMREAD_GRAYSCALE)).astype(np.float32) lm_img = cv2.resize(lm_img, (RESIZE_HEIGHT, RESIZE_WIDTH)) lm_img = np.expand_dims(lm_img, 2) else: lm_img = np.zeros((RESIZE_HEIGHT, RESIZE_WIDTH, 1)) # 把landmark拼接上 buffer[i] = np.concatenate((frame, lm_img),2) else: buffer[i] = frame buffer = (buffer - 128) / 128 buffer = buffer.transpose((3, 0, 1, 2)) inputs = torch.from_numpy(buffer).unsqueeze(0) # if not split == 'Submission_Set': # avgface_path = frame_root_dir.replace('image', 'avgfaces') # avgface_path = os.path.join(avgface_path, 'avgface.png') # avgface = cv2.imread(avgface_path, cv2.IMREAD_GRAYSCALE) # avgface = np.stack((avgface,)*3, axis=-1) # avgface = (avgface - 128) / 128 # avgface = torch.from_numpy(avgface.transpose(2,0,1)).float().unsqueeze(0) if not split == 'Submission_Set': # 取label with open(label_path[0], 'rt') as f: lines = f.read().splitlines() line = lines[center] #去掉第一行 valence_label[center-1, 0] = float(line.split(',')[0]) arousal_label[center-1, 0] = float(line.split(',')[1]) with torch.no_grad(): # 方式1 outputs_va, outputs_expr, _, _ = model(inputs.cuda()) outputs_v = outputs_va[:,:20] outputs_a = outputs_va[:,20:40] _, valence_pred_reg, arousal_pred_reg = criterion(outputs_v, outputs_a, \ valence_label, arousal_label) # 方式2 # logits1, logits2, logits3, corr_t_div_C, corr_s_div_C, feat = model(inputs.cuda()) # outputs_v2, outputs_a2 = logits2[0][:,:20], logits2[0][:,20:40] # _, valence_pred_reg2, arousal_pred_reg2 = criterion(outputs_v2, outputs_a2, \ # valence_label, arousal_label) # valence_pred_reg = valence_pred_reg2 # arousal_pred_reg = arousal_pred_reg2 # visualize_feature(feat, center) # visualize_attention(inputs, corr_t_div_C, corr_s_div_C) pred_concat = np.concatenate((np.expand_dims(valence_pred_reg,0), np.expand_dims(arousal_pred_reg,0)), 1) preds_all[center-1, :] = pred_concat # 将无效label去掉 if valence_label[center-1, 0] == -5: preds_all[center-1, :] = -5 preds_all_valid = [] valence_label_valid = [] arousal_label_valid = [] for each in preds_all: if not each[0] == -5: preds_all_valid.append(each) for each in valence_label: if not each[0] == -5: valence_label_valid.append(each) for each in arousal_label: if not each[0] == -5: arousal_label_valid.append(each) print(len(preds_all_valid), len(valence_label_valid), len(arousal_label_valid)) preds_all = torch.from_numpy(np.array(preds_all_valid)) valence_label = torch.from_numpy(np.array(valence_label_valid)) arousal_label = torch.from_numpy(np.array(arousal_label_valid)) val_cc2 = concord_cc2(preds_all[:,0] , valence_label[:,0]) aro_cc2 = concord_cc2(preds_all[:,1], arousal_label[:,0]) val_mse = torch.nn.MSELoss()(preds_all[:,0], valence_label[:,0]) aro_mse = torch.nn.MSELoss()(preds_all[:,1], arousal_label[:,0]) if not split == 'Submission_Set': if not os.path.exists(os.path.join(save_dir, 'res')): os.mkdir(os.path.join(save_dir, 'res')) with open(os.path.join(save_dir, 'res', 'res_' + frames[0][0].split('/')[-2] + '.txt'), 'w') as res: for i, each in enumerate(preds_all): res.write(str(preds_all.numpy()[i,0]) + ", " + str(valence_label.numpy()[i,0]) + ', ' \ + str(preds_all.numpy()[i,1]) + ", " + str(arousal_label.numpy()[i,0]) + '\n') if split == 'Submission_Set': if not os.path.exists(os.path.join(save_dir, 'submission')): os.mkdir(os.path.join(save_dir, 'submission')) with open(os.path.join(save_dir, 'submission', frames[0][0].split('/')[-2] + '.txt'), 'w') as res: res.write("valence,arousal\n") for i, each in enumerate(preds_all): res.write(str(preds_all.numpy()[i,0]) + "," + str(preds_all.numpy()[i,1]) + '\n') print("Val CCC: {:.4f} Aro CCC: {:.4f} Val MSE: {:.4f} Aro MSE: {:.4f}".format(val_cc2, aro_cc2, val_mse, aro_mse)) with open('test_log.txt', 'a') as test_log: log = frames[0][0].split('/')[-2] + ',' + str(float(val_cc2)) + ',' + str(float(aro_cc2)) + ',' + \ str(float(val_mse)) + ',' + str(float(aro_mse)) test_log.write(log + '\n') val_cccs.append(val_cc2) aro_cccs.append(aro_cc2) val_mses.append(val_mse) aro_mses.append(aro_mse) print("[test] Val CCC: {:.4f} Aro CCC: {:.4f} Val MSE: {:.4f} Aro MSE: {:.4f}".format(np.mean(val_cccs), \ np.mean(aro_cccs), np.mean(val_mses), np.mean(aro_mses))) with open('test_log.txt', 'a') as test_log: test_log.write(("[test] Val CCC: {:.4f} Aro CCC: {:.4f} Val MSE: {:.4f} Aro MSE: {:.4f}\n".format(np.mean(val_cccs), \ np.mean(aro_cccs), np.mean(val_mses), np.mean(aro_mses)))) stop_time = timeit.default_timer() print("Execution time: " + str(stop_time - start_time) + "\n")