def __init__(self, teacher_path):
super().__init__()
teacher_yaml = teacher_path.parent / 'config.yaml'
teacher_args = load_yaml(teacher_yaml)['model_args']
teacher = Network(**teacher_args)
teacher.load_state_dict(torch.load(str(teacher_path)))
self.teacher = teacher
self.converter = ConverterTorch()
def main(config):
# from thop import profile
# net = Network(**config['model_args'])
# input = torch.randn(1, 6, 160, 384) # 模型输入的形状,batch_size=1
# flops, params = profile(net, inputs=(input,))
# print(flops / 1e9, 'G', params / 1e6) # flops单位G，para单位M
# aa
data_train, data_val = get_dataset(config['source'])(**config['data_args'])
net = Network(**config['model_args']).to(config['device'])
net.load_state_dict(torch.load(str(config['model_path'])))
# wandb.init(
# project='task-distillation-eval',
# config=config, id=config['run_name'], resume='auto')
# wandb.save(str(Path(wandb.run.dir) / '*.t7'))
with torch.no_grad():
MAE_val, RMSE_val, EVS_val = net_eval(net, data_val, config)
print(' MAE_val', MAE_val)
print(' RMSE_val: ', RMSE_val)
print(' EVS_val: ', EVS_val)
class Planner(object):
def __init__(self, path_to_conf_file):
self.device = torch.device(
'cuda' if torch.cuda.is_available() else 'cpu')
self.transform = torchvision.transforms.ToTensor()
self.converter = ConverterTorch().to(self.device)
self.target_index = 65
self.speed_mult = 2.5
path_to_conf_file = Path(path_to_conf_file)
config = load_yaml(path_to_conf_file.parent / 'config.yaml')
self.net = Network(**config['model_args']).to(self.device)
self.net.load_state_dict(torch.load(path_to_conf_file))
self.net.eval()
@torch.no_grad()
def run_step(self, rgb, rgb_forward, viz=None):
if Modular:
# Modularity and Abstract
rgb = Image.fromarray(rgb).convert('RGB')
img = input_transform_cityscapes(rgb)
img = img.cuda().unsqueeze(0)
rgb_forward = Image.fromarray(rgb_forward).convert('RGB')
img_forward = input_transform_cityscapes(rgb_forward)
img_forward = img_forward.cuda().unsqueeze(0)
output = model(img)
label = output[0].max(0)[1].byte().cpu().data
label_color = Colorize()(label.unsqueeze(0))
rgb = ToPILImage()(label_color)
rgb.save('./seg.jpg')
output = model(img_forward)
label = output[0].max(0)[1].byte().cpu().data
label_color = Colorize()(label.unsqueeze(0))
rgb_forward = ToPILImage()(label_color)
rgb_forward.save('./seg_2.jpg')
img = self.transform(rgb).to(self.device).unsqueeze(0)
img_forward = self.transform(rgb_forward).to(self.device).unsqueeze(0)
# print(img_forward.shape)
model_input = torch.cat((img_forward, img), 1)
cam_coords = self.net(model_input)
cam_coords[..., 0] = (cam_coords[..., 0] +
1) / 2 * img.shape[-1] # rgb coords
cam_coords[..., 1] = (cam_coords[..., 1] + 1) / 2 * img.shape[-2]
map_coords = self.converter.cam_to_map(
cam_coords).cpu().numpy().squeeze()
world_coords = self.converter.cam_to_world(
cam_coords).cpu().numpy().squeeze()
target_speed = np.sqrt(
((world_coords[:2] - world_coords[1:3])**2).sum(1).mean())
target_speed *= self.speed_mult
theta1 = np.degrees(np.arctan2(world_coords[0][0], world_coords[0][1]))
theta2 = np.degrees(np.arctan2(world_coords[4][0], world_coords[4][1]))
# print(abs(theta2 - theta1))
if abs(theta2 - theta1) < 2:
target_speed *= self.speed_mult
else:
target_speed *= 1.2
curve = spline(map_coords + 1e-8 * np.random.rand(*map_coords.shape),
100)
target = curve[self.target_index]
curve_world = spline(
world_coords + 1e-8 * np.random.rand(*world_coords.shape), 100)
target_world = curve_world[self.target_index]
if viz:
viz.planner_draw(cam_coords.cpu().numpy().squeeze(), map_coords,
curve, target)
return target_world, target_speed
class DQNAgent:
def __init__(
self,
env: UnityEnvironment,
memory_size: int,
batch_size: int,
target_update: int,
epsilon_decay: float = 1 / 2000,
max_epsilon: float = 1.0,
min_epsilon: float = 0.1,
gamma: float = 0.99,
):
self.brain_name = env.brain_names[0]
self.brain = env.brains[self.brain_name]
env_info = env.reset(train_mode=True)[self.brain_name]
self.env = env
action_size = self.brain.vector_action_space_size
state = env_info.vector_observations[0]
state_size = len(state)
self.obs_dim = state_size
self.action_dim = 1
self.memory = ReplayBuffer(self.obs_dim, self.action_dim, memory_size, batch_size)
self.batch_size = batch_size
self.target_update = target_update
self.epsilon_decay = epsilon_decay
self.max_epsilon = max_epsilon
self.min_epsilon = min_epsilon
self.gamma = gamma
self.epsilon = max_epsilon
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.dqn = Network(self.obs_dim, self.action_dim)
self.dqn_target = Network(self.obs_dim, self.action_dim)
self.dqn_target.load_state_dict(self.dqn.state_dict())
self.dqn_target.eval()
self.optimizer = optim.Adam(self.dqn.parameters(), lr=5e-5)
self.transition = list()
self.is_test = False
def select_action(self, state: np.ndarray) -> np.int64:
""" Select an action given input """
if self.epsilon > np.random.random():
selected_action = np.random.random_integers(0, self.action_dim-1)
else:
selected_action = self.dqn(
torch.FloatTensor(state).to(self.device)
)
selected_action = np.argmax(selected_action.detach().cpu().numpy())
if not self.is_test:
self.transition = [state, selected_action]
return selected_action
def step(self, action: np.int64) -> Tuple[np.ndarray, np.float64, bool]:
"Take an action and return environment response"
env_info = self.env.step(action)[self.brain_name]
next_state = env_info.vector_observations[0]
reward = env_info.rewards[0]
done = env_info.local_done[0]
if not self.is_test:
self.transition += [reward, next_state, done]
self.memory.store(*self.transition)
return next_state, reward, done
def update_model(self) -> torch.Tensor:
""" Update model by gradient descent"""
samples = self.memory.sample_batch()
loss = self._compute_dqn_loss(samples)
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
return loss.item()
def train(self, num_episode: int, max_iteration: int=1000, plotting_interval: int=400):
""" train the agent """
self.is_test = False
env_info = self.env.reset(train_mode=True)[self.brain_name]
state = env_info.vector_observations[0]
update_cnt = 0
epsilons = []
losses = []
avg_losses= []
scores = []
avg_scores = []
for episode in range(num_episode):
env_info = self.env.reset(train_mode=True)[self.brain_name]
state = env_info.vector_observations[0]
score = 0
for iter in range(max_iteration):
action = self.select_action(state)
next_state, reward, done = self.step(action)
state = next_state
score += reward
if done:
break
if len(self.memory) > self.batch_size:
loss = self.update_model()
losses.append(loss)
update_cnt += 1
avg_losses.append(np.mean(losses))
losses = []
self.epsilon = max(
self.min_epsilon, self.epsilon - (
self.max_epsilon - self.min_epsilon
) * self.epsilon_decay
)
epsilons.append(self.epsilon)
if update_cnt % self.target_update == 0:
self._target_hard_update()
scores.append(score)
epsilons.append(self.epsilon)
if episode >= 100:
avg_scores.append(np.mean(scores[-100:]))
self._plot(episode, scores, avg_scores, avg_losses, epsilons)
torch.save(self.dqn.state_dict(), "model_weight/dqn.pt")
def test(self):
""" Test agent """
self.is_test = True
env_info = self.env.reset(train_mode=False)[self.brain_name]
state = env_info.vector_observations[0]
done = False
score = 0
while not done:
action = self.select_action(state)
next_state, reward, done = self.step(action)
state = next_state
score += reward
print("score: ", score)
self.env.close()
def _compute_dqn_loss(self, samples: Dict[str, np.ndarray], gamma: float=0.99) -> torch.Tensor:
""" Compute and return DQN loss"""
gamma = self.gamma
device = self.device
state = torch.FloatTensor(samples["obs"]).to(device)
next_state = torch.FloatTensor(samples["next_obs"]).to(device)
action = torch.LongTensor(samples["acts"]).reshape(-1, 1).to(device)
reward = torch.FloatTensor(samples["rews"]).reshape(-1, 1).to(device)
done = torch.FloatTensor(samples["done"]).reshape(-1, 1).to(device)
curr_q_value = self.dqn(state).gather(1, action)
next_q_value = self.dqn_target(next_state).max(dim=1, keepdim=True)[0].detach()
mask = 1 - done
target = (reward + gamma * next_q_value * mask).to(device)
loss = F.smooth_l1_loss(curr_q_value, target)
return loss
def _target_hard_update(self):
""" update target network """
self.dqn_target.load_state_dict(self.dqn.state_dict())
def _plot(
self,
episode :int,
scores: List[float],
avg_scores: List[float],
losses: List[float],
epsilons: List[float]
):
""" Plot the training process"""
plt.figure(figsize=(20, 5))
plt.subplot(141)
if len(avg_scores) > 0:
plt.title("Average reward per 100 episodes. Score: %s" % (avg_scores[-1]))
else:
plt.title("Average reward over 100 episodes.")
plt.plot([100 + i for i in range(len(avg_scores))], avg_scores)
plt.subplot(142)
plt.title("episode %s. Score: %s" % (episode, np.mean(scores[-10:])))
plt.plot(scores)
plt.subplot(143)
plt.title('Loss')
plt.plot(losses)
plt.subplot(144)
plt.title('epsilons')
plt.plot(epsilons)
plt.savefig('plots/dqn_result.png')
if __name__ == "__main__":
args = parse_args()
coloredlogs.install(
level="INFO", fmt="%(asctime)s %(filename)s %(levelname)s %(message)s")
logging.info("Called with args: " + str(args))
if args.cfg:
cfg_from_file(args.cfg)
net = Network()
checkpoint = torch.load(osp.abspath(args.checkpoint))
net.load_state_dict(checkpoint["model"])
logging.info("Loaded checkpoint from: %s" % args.checkpoint)
net.eval()
device = torch.device("cuda:%s" % args.gpu if args.gpu != -1 else "cpu")
net.to(device)
# Extract feature of the query person
query_img = cv2.imread("imgs/query.jpg")
query_roi = np.array([0, 0, 466, 943]) # [x1, y1, x2, y2]
query_feat = net.inference(query_img, query_roi).view(-1, 1)
# Get gallery images
gallery_imgs = sorted(glob("imgs/gallery*.jpg"))
for gallery_img in gallery_imgs:
logging.info("Detecting %s" % gallery_img)
output_dir = osp.join(cfg.DATA_DIR, "trained_model")
if not os.path.exists(output_dir):
os.makedirs(output_dir)
assert args.dataset in ["psdb_train",
"psdb_test"], "Unknown dataset: %s" % args.dataset
dataset = PSDB(args.dataset)
dataloader = DataLoader(dataset, batch_size=1, sampler=PSSampler(dataset))
logging.info("Loaded dataset: %s" % args.dataset)
# Initialize model
net = Network()
if args.weights:
state_dict = torch.load(args.weights)
net.load_state_dict(
{k: v
for k, v in state_dict.items() if k in net.state_dict()})
logging.info("Loaded pretrained model from: %s" % args.weights)
# Initialize optimizer
lr = cfg.TRAIN.LEARNING_RATE
weight_decay = cfg.TRAIN.WEIGHT_DECAY
params = []
for k, v in net.named_parameters():
if v.requires_grad:
if "BN" in k:
params += [{"params": [v], "lr": lr, "weight_decay": 0}]
elif "bias" in k:
params += [{"params": [v], "lr": 2 * lr, "weight_decay": 0}]
else:
params += [{
def main():
if not torch.cuda.is_available():
print('No GPU device available')
sys.exit(1)
np.random.seed(args.seed)
cudnn.benchmark = True
torch.manual_seed(args.seed)
cudnn.enabled = True
torch.cuda.manual_seed(args.seed)
print("args = %s", args)
print("unparsed args = %s", unparsed)
# prepare dataset
if args.cifar100:
train_transform, valid_transform = utils._data_transforms_cifar100(args)
else:
train_transform, valid_transform = utils._data_transforms_cifar10(args)
if args.cifar100:
train_data = dset.CIFAR100(root=args.tmp_data_dir, train=True, download=True, transform=train_transform)
valid_data = dset.CIFAR100(root=args.tmp_data_dir, train=False, download=True, transform=valid_transform)
else:
train_data = dset.CIFAR10(root=args.tmp_data_dir, train=True, download=True, transform=train_transform)
valid_data = dset.CIFAR10(root=args.tmp_data_dir, train=False, download=True, transform=valid_transform)
train_queue = torch.utils.data.DataLoader(
train_data, batch_size=args.batch_size, shuffle=True, pin_memory=True, num_workers=args.workers)
valid_queue = torch.utils.data.DataLoader(
valid_data, batch_size=args.batch_size, shuffle=False, pin_memory=True, num_workers=args.workers)
ood_queues = {}
for k in ['svhn', 'lsun_resized', 'imnet_resized']:
ood_path = os.path.join(args.ood_dir, k)
dset_ = dset.ImageFolder(ood_path, valid_transform)
loader = torch.utils.data.DataLoader(
dset_, batch_size=args.batch_size, shuffle=False,
pin_memory=True, num_workers=args.workers
)
ood_queues[k] = loader
# build Network
criterion = nn.CrossEntropyLoss()
criterion = criterion.cuda()
supernet = Network(
args.init_channels, CIFAR_CLASSES, args.layers,
combine_method=args.feat_comb, is_cosine=args.is_cosine,
)
supernet.cuda()
# print(len(supernet.cells))
ckpt = torch.load(args.load_at)
print(args.load_at)
supernet.load_state_dict(ckpt)
supernet.generate_share_alphas()
# alphas = torch.Tensor([
# [0., 1., 1.],
# [0., 1., 0.],
# [0., 1., 0.],
# [0., 1., 1.],
# [0., 1., 1.],
# [0., 1., 1.],
# [0., 1., 0.],
# [0., 1., 0.],
# [0., 1., 1.],
# [0., 1., 0.],
# [0., 1., 0.],
# [0., 1., 1.],
# [0., 1., 0.],
# [0., 1., 1.]
# ]).cuda()
# for i in range(8):
# supernet.cells[i].ops_alphas = alphas
alphas = supernet.cells[0].ops_alphas
print(alphas)
out_dir = './results/{}/eval_out/{}'.format(args.load_at.split('/')[2], args.seed)
if not os.path.exists(out_dir):
os.makedirs(out_dir)
torch.save(alphas, os.path.join(out_dir, 'alphas.pt'))
with open(os.path.join(out_dir, 'alphas.txt'), 'w') as f:
for i in alphas.cpu().detach().numpy():
for j in i:
f.write('{:d}'.format(int(j)))
f.write('\n')
if args.cifar100:
weight_decay = 5e-4
else:
weight_decay = 3e-4
optimizer = torch.optim.SGD(
supernet.parameters(),
args.learning_rate,
momentum=args.momentum,
weight_decay=weight_decay,
)
# scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, float(args.epochs), eta_min=0)
valid_acc, _ = infer(valid_queue, supernet, criterion)
print('valid_acc {:.2f}'.format(valid_acc))
lg_aucs, sm_aucs, ent_aucs = ood_eval(valid_queue, ood_queues, supernet, criterion)
with open(os.path.join(out_dir, 'before.txt'), 'w') as f:
f.write('-'.join([str(valid_acc), str(lg_aucs), str(sm_aucs), str(ent_aucs)]))
if args.fine_tune:
for epoch in range(args.epochs):
# scheduler.step()
print('epoch {} lr {:.4f}'.format(epoch, 0.001))#scheduler.get_lr()[0]))
train_acc, _ = train(train_queue, supernet, criterion, optimizer)
print('train_acc {:.2f}'.format(train_acc))
valid_acc, _ = infer(valid_queue, supernet, criterion)
print('valid_acc {:.2f}'.format(valid_acc))
lg_aucs, sm_aucs, ent_aucs = ood_eval(valid_queue, ood_queues, supernet, criterion)
with open(os.path.join(out_dir, 'after.txt'), 'w') as f:
f.write('-'.join([str(valid_acc), str(lg_aucs), str(sm_aucs), str(ent_aucs)]))
def main():
if not torch.cuda.is_available():
print('No GPU device available')
sys.exit(1)
np.random.seed(args.seed)
cudnn.benchmark = True
torch.manual_seed(args.seed)
cudnn.enabled = True
torch.cuda.manual_seed(args.seed)
print("args = %s", args)
print("unparsed args = %s", unparsed)
# prepare dataset
if args.cifar100:
train_transform, valid_transform = utils._data_transforms_cifar100(args)
else:
train_transform, valid_transform = utils._data_transforms_cifar10(args)
if args.cifar100:
train_data = dset.CIFAR100(root=args.tmp_data_dir, train=True, download=False, transform=train_transform)
valid_data = dset.CIFAR100(root=args.tmp_data_dir, train=False, download=False, transform=valid_transform)
else:
train_data = dset.CIFAR10(root=args.tmp_data_dir, train=True, download=False, transform=train_transform)
valid_data = dset.CIFAR10(root=args.tmp_data_dir, train=False, download=False, transform=valid_transform)
train_queue = torch.utils.data.DataLoader(
train_data, batch_size=args.batch_size, shuffle=True, pin_memory=True, num_workers=args.workers)
valid_queue = torch.utils.data.DataLoader(
valid_data, batch_size=args.batch_size, shuffle=False, pin_memory=True, num_workers=args.workers)
# build Network
criterion = nn.CrossEntropyLoss()
criterion = criterion.cuda()
supernet = Network(
args.init_channels, CIFAR_CLASSES, args.layers
)
supernet.cuda()
ckpt = torch.load(args.load_at)
print(args.load_at)
supernet.load_state_dict(ckpt)
supernet.generate_share_alphas()
alphas = supernet.cells[0].ops_alphas
print(alphas)
out_dir = args.save + '{}/eval_out/{}'.format(args.load_at.split('/')[2], args.seed)
if not os.path.exists(out_dir):
os.makedirs(out_dir)
torch.save(alphas, os.path.join(out_dir, 'alphas.pt'))
with open(os.path.join(out_dir, 'alphas.txt'), 'w') as f:
for i in alphas.cpu().detach().numpy():
for j in i:
f.write('{:d}'.format(int(j)))
f.write('\n')
# Getting subnet according to sample alpha
subnet = supernet.get_sub_net(alphas)
init_valid_acc, _ = infer(valid_queue, subnet, criterion)
print('Initial Valid Acc {:.2f}'.format(init_valid_acc))
if args.fine_tune:
if args.cifar100:
weight_decay = 5e-4
else:
weight_decay = 3e-4
# Fine tuning whole network:
subnet = supernet.get_sub_net(alphas)
optimizer = torch.optim.SGD(
subnet.parameters(),
args.finetune_lr,
momentum=args.momentum,
weight_decay=weight_decay,
)
for epoch in range(args.epochs):
# scheduler.step()
print('epoch {} lr {:.4f}'.format(epoch, args.finetune_lr))
train_acc, _ = train(train_queue, subnet, criterion, optimizer)
print('train_acc {:.2f}'.format(train_acc))
whole_valid_acc, _ = infer(valid_queue, subnet, criterion)
print('valid_acc after whole fine-tune {:.2f}'.format(whole_valid_acc))
fly_whole_valid_acc, _ = infer(valid_queue, subnet, criterion, use_fly_bn=False)
print('valid_acc after whole fine-tune {:.2f}'.format(fly_whole_valid_acc))
# Fine-tuning only classifier:
subnet = supernet.get_sub_net(alphas)
# Freezing other weights except classifier:
for name, param in subnet.named_parameters():
if not 'classifier' in name:
param.requires_grad_(requires_grad=False)
optimizer = torch.optim.SGD(
subnet.classifier.parameters(),
args.finetune_lr,
momentum=args.momentum,
weight_decay=weight_decay,
)
for epoch in range(args.epochs):
# scheduler.step()
print('epoch {} lr {:.4f}'.format(epoch, args.finetune_lr))
train_acc, _ = train(train_queue, subnet, criterion, optimizer)
print('train_acc {:.2f}'.format(train_acc))
part_valid_acc, _ = infer(valid_queue, subnet, criterion)
print('valid_acc after fine-tuning classifier {:.2f}'.format(part_valid_acc))
fly_part_valid_acc, _ = infer(valid_queue, subnet, criterion, use_fly_bn=False)
print('valid_acc after fine-tuning classifier {:.2f}'.format(fly_part_valid_acc))
with open(os.path.join(out_dir, 'results.txt'), 'w') as f:
f.write('-'.join([str(init_valid_acc), str(whole_valid_acc),
str(fly_whole_valid_acc), str(part_valid_acc), str(fly_part_valid_acc)]))
if not args.fine_tune:
with open(os.path.join(out_dir, 'results.txt'), 'w') as f:
f.write(str(init_valid_acc))
