parser = argparse.ArgumentParser(description="attribute detection") parser.add_argument("--trained_model", type=str) # parser.add_argument("--label_file", type=str, help="The label file path.") # parser.add_argument("--nms_method", type=str, default="hard") args = parser.parse_args() model_kwargs = {'drop_pool5': True, 'drop_pool5_rate': 0.5, 'last_conv_stride': 2, 'num_att': 35} DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") if __name__ == '__main__': model = DeepMAR_ResNet50(**model_kwargs) state = torch.load('/home/thulab/Desktop/pedestrian-attribute-recognition-pytorch/exp/deepmar_resnet50/pa100k/partition0/run1/model/ckpt_epoch150.pth') # print(type(state['state_dicts'])) # print(len(state['state_dicts'])) # print(state['state_dicts'][0]) model.load_state_dict(state['state_dicts'][0], strict=True) example = torch.rand(1, 3, 224, 224) torch_out = torch.onnx.export(model, example, "test.onnx", verbose=True ) print('onnx done!')
# set the random seed if cfg.set_seed: set_seed( cfg.rand_seed ) # init the gpu ids set_devices(cfg.sys_device_ids) # dataset normalize = transforms.Normalize(mean=cfg.mean, std=cfg.std) test_transform = transforms.Compose([ transforms.Resize(cfg.resize), transforms.ToTensor(), normalize,]) ### Att model ### model = DeepMAR_ResNet50(**cfg.model_kwargs) # load model weight if necessary if cfg.load_model_weight: map_location = (lambda storage, loc:storage) ckpt = torch.load(cfg.model_weight_file, map_location=map_location) model.load_state_dict(ckpt['state_dicts'][0]) model.cuda() model.eval() # load one image img = Image.open(cfg.demo_image) img_trans = test_transform( img ) img_trans = torch.unsqueeze(img_trans, dim=0) img_var = Variable(img_trans).cuda()
num_workers=cfg.workers, pin_memory=True, drop_last=False) test_transform = transforms.Compose([ transforms.Resize(cfg.resize), transforms.ToTensor(), normalize, ]) test_set = AttDataset(dataset=cfg.dataset, partition=cfg.partition, split=cfg.test_split, partition_idx=cfg.partition_idx, transform=test_transform) ### Att model ### model = DeepMAR_ResNet50(**cfg.model_kwargs) # 使用ResNet-50进行多标签分类 # Wrap the model after set_devices, data parallel model_w = torch.nn.DataParallel(model) # 使用DataParallel进行单机多卡训练 # using the weighted cross entropy loss if cfg.weighted_entropy: rate = np.array(train_set.partition['weight_' + cfg.split][cfg.partition_idx]) rate = rate[train_set.dataset['selected_attribute']].tolist() else: rate = None # compute the weight of positive and negative if rate is None: weight_pos = [1 for i in range(num_att)] weight_neg = [1 for i in range(num_att)]