def train_model(model, criterion1, criterion2, optimizer, scheduler, num_epochs, use_gpu, batchnumber=50, batchsize=32): since = time.time() writer = SummaryWriter() d = get_imgiddic( '/home/csc302/bishe/dataset/VehicleID_V1.0/train_50000.txt') dsort = get_img_sortiddic( '/home/csc302/bishe/dataset/VehicleID_V1.0/train_50000.txt') for epoch in range(num_epochs): begin_time = time.time() print('Epoch {}/{}'.format(epoch, num_epochs - 1)) print('-' * 10) scheduler.step() model.train(True) # Set model to training mode running_loss = 0.0 running_corrects = 0 running_softmaxloss = 0.0 running_msml = 0.0 for batch in range(batchnumber): imglist, idlist = gettriphardsample( '/home/csc302/bishe/dataset/VehicleID_V1.0/triplet_train/', batchsize=batchsize, imgiddic=d, imgsortid=dsort) image_datasetsoft = Mydatsetsoft(imglist, idlist, data_transforms=data_transforms) data_loadsoft = torch.utils.data.DataLoader(image_datasetsoft, batch_size=batchsize * 4, shuffle=False) for data in data_loadsoft: inputs, labels = data # wrap them in Variable if use_gpu: inputs = Variable(inputs.cuda()) labels = Variable(labels.cuda()) else: inputs, labels = Variable(inputs), Variable(labels) optimizer.zero_grad() # forward feature, outputs = model(inputs) _, preds = torch.max(outputs.data, 1) loss1 = criterion1(outputs, labels) running_softmaxloss += loss1.data[0] #feature_p, feature_n = maketrihardbatch(feature) feature = normalize(feature) loss2 = maketrihardbatch(feature, labels, criterion2) running_msml += loss2.data[0] loss = loss1 + loss2 #embed() loss.backward() optimizer.step() running_loss += loss.data[0] running_corrects += torch.sum(preds == labels.data) #optimizer.zero_grad() # forward # loss.backward() # optimizer.step() # running_loss += loss.data[0] #embed() # print result every 10 batch if (batch + 1) % 30 == 0: batch_loss = running_loss / (batch + 1) batch_softmaxloss = running_softmaxloss / (batch + 1) batch_msml = running_msml / (batch + 1) batch_acc = running_corrects / (4 * (batch + 1) * batchsize) print('Epoch [{}] Batch [{}] Loss: {:.4f} SoftmaxLoss: {:.4f} TriphardLoss:{:.4f} Accuracy:{:.4f} Time: {:.4f}s'. \ format(epoch, batch+1, batch_loss,batch_softmaxloss, batch_msml, batch_acc, time.time()-begin_time)) begin_time = time.time() epoch_loss = running_loss / batchnumber epoch_acc = running_corrects / (4 * batchnumber * batchsize) epoch_softmaxloss = running_softmaxloss / batchnumber epoch_msml = running_msml / batchnumber print( 'Loss: {:.4f} SoftmaxLoss: {:.4f} TriphardLoss: {:.4f} Accuracy:{:.4f}' .format(epoch_loss, epoch_softmaxloss, epoch_msml, epoch_acc)) if not os.path.exists('output'): os.makedirs('output') torch.save(model, 'output/resnet_epoch{}.pkl'.format(epoch)) if (epoch + 1) % 10 == 0: model_wtse = model.state_dict() model_nofce = resnet50_nofc(pretrained=False) model_nofce.load_state_dict(remove_fc(model_wtse)) gallery, probe, gdict, pdict = get_galproset( '/home/csc302/bishe/dataset/VehicleID_V1.0/train_test_split/test_list_800.txt' ) gallerydict, probedict = getfeature( imgpath='/home/csc302/bishe/dataset/VehicleID_V1.0/test_800/', model=model_nofce.cuda(), gallery=gallery, probe=probe) print( calacc(gallerydict=gallerydict, probedict=probedict, gdict=gdict, pdict=pdict)) writer.add_scalar('epoch_loss', epoch_loss, epoch) writer.add_scalar('epoch_softmax_loss', epoch_softmaxloss, epoch) writer.add_scalar('epoch_triphard_loss', epoch_msml, epoch) time_elapsed = time.time() - since print('Training complete in {:.0f}m {:.0f}s'.format( time_elapsed // 60, time_elapsed % 60)) writer.export_scalars_to_json("./all_scalars.json") writer.close() return model
def train_model(model, criterion, optimizer, scheduler, num_epochs, use_gpu): since = time.time() writer = SummaryWriter() best_model_wts = model.state_dict() best_acc = 0.0 for epoch in range(num_epochs): begin_time = time.time() count_batch = 0 print('Epoch {}/{}'.format(epoch, num_epochs - 1)) print('-' * 10) # Each epoch has a training and validation phase for phase in ['train', 'val']: if phase == 'train': scheduler.step() model.train(True) # Set model to training mode else: model.train(False) # Set model to evaluate mode running_loss = 0.0 running_corrects = 0 # Iterate over data. for data in dataloders[phase]: count_batch += 1 # get the inputs inputs, labels = data # wrap them in Variable if use_gpu: inputs = Variable(inputs.cuda()) labels = Variable(labels.cuda()) else: inputs, labels = Variable(inputs), Variable(labels) # zero the parameter gradients optimizer.zero_grad() # forward outputs = model(inputs) _, preds = torch.max(outputs.data, 1) loss = criterion(outputs, labels) embed() # backward + optimize only if in training phase if phase == 'train': loss.backward() optimizer.step() # statistics running_loss += loss.data[0] running_corrects += torch.sum(preds == labels.data) # print result every 10 batch if count_batch % 10 == 0: batch_loss = running_loss / (batch_size * count_batch) batch_acc = running_corrects / (batch_size * count_batch) print('{} Epoch [{}] Batch [{}] Loss: {:.4f} Acc: {:.4f} Time: {:.4f}s'. \ format(phase, epoch, count_batch, batch_loss, batch_acc, time.time()-begin_time)) begin_time = time.time() epoch_loss = running_loss / dataset_sizes[phase] epoch_acc = running_corrects / dataset_sizes[phase] #print(running_corrects, running_loss, dataset_sizes[phase]) print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc)) # save model if phase == 'train': if not os.path.exists('output'): os.makedirs('output') torch.save(model, 'output/resnet_epoch{}.pkl'.format(epoch)) writer.add_scalar('train_epoch_loss', epoch_loss, epoch) writer.add_scalar('train_epoch_acc', epoch_acc, epoch) if phase == 'val': writer.add_scalar('val_epoch_loss', epoch_loss, epoch) writer.add_scalar('val_epoch_acc', epoch_acc, epoch) # deep copy the model if phase == 'val' and epoch_acc > best_acc: best_acc = epoch_acc best_model_wts = model.state_dict() time_elapsed = time.time() - since print('Training complete in {:.0f}m {:.0f}s'.format( time_elapsed // 60, time_elapsed % 60)) print('Best val Acc: {:4f}'.format(best_acc)) writer.export_scalars_to_json("./all_scalars.json") writer.close() # load best model weights model.load_state_dict(best_model_wts) feature_model = resnet50_nofc(pretrained=False) #embed() feature_model.load_state_dict(remove_fc(best_model_wts)) return model, feature_model
# f1 = open('gdict.txt', 'wb') # pickle.dump(featDict, f1) # f1 = open('gdict.txt', 'rb') # featDict = pickle.load(f1) # f1.close() for i in range(test_num): probeDict, galleryDict = getGalleryProbe(imageInfo) cmc_result.append( calacc(featDict, probeDict, galleryDict, norm_flag, rerank)) print('test time: ', i) cmc_mean = np.mean(np.array(cmc_result), axis=0) return cmc_mean if __name__ == "__main__": imageInfo, imageList = getAllImamgeInfo( '/home/csc302/bishe/dataset/VehicleID_V1.0/train_test_split/test_list_800.txt' ) modelfc = torch.load( '/home/csc302/bishe/代码/VehicleReID/output-triphard+softmax104/resnet_nofc_epoch104.pkl' ) wts = modelfc.state_dict() model = resnet50_nofc(pretrained=False) model.load_state_dict(remove_fc(wts)) model = model.cuda() cmc_mean = cmc(model, test_num=10, norm_flag=True, rerank=False) print(cmc_mean)