def get_classifier(ebd_dim, args):
tprint("Building classifier")
model = MLP(ebd_dim, args)
if args.cuda != -1:
return model.cuda(args.cuda)
else:
return model
def main():
# get unity environment
env, brain = get_unity_envs()
# get arguments
args = get_arguments()
print(args)
# set gpu environment
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
cudnn.enabled = True
cudnn.benchmark = True
cuda = torch.cuda.is_available()
# set random seed
rn = set_seeds(args.random_seed, cuda)
# make directory
os.makedirs(args.snapshot_dir, exist_ok=True)
# get validation dataset
val_set = get_validation_dataset(args)
print("len of test set: ", len(val_set))
val_loader = data.DataLoader(val_set,
batch_size=args.real_batch_size,
shuffle=False,
num_workers=args.num_workers,
pin_memory=True)
# generate training list
with open(args.syn_list_path, "w") as fp:
for i in range(args.syn_img_num):
if i % 10 != 0:
fp.write(str(i + 1) + '\n')
# get main model
main_model = MLP(args.num_inputs, args.num_outputs, args.hidden_size)
if args.resume != "":
main_model.load_state_dict(torch.load(args.resume))
# get task model
if args.task_model_name == "FCN8s":
task_model = FCN8s_sourceonly(n_class=args.num_classes)
vgg16 = VGG16(pretrained=True)
task_model.copy_params_from_vgg16(vgg16)
else:
raise ValueError("Specified model name: FCN8s")
# save initial task model
torch.save(task_model.state_dict(),
os.path.join(args.snapshot_dir, "task_model_init.pth"))
if cuda:
main_model = main_model.cuda()
task_model = task_model.cuda()
# get optimizer
main_optimizer = optim.Adam(main_model.parameters(), lr=args.main_lr)
task_optimizer = optim.SGD(task_model.parameters(),
lr=args.task_lr,
momentum=0.9,
weight_decay=1e-4)
frame_idx = 0
whole_start_time = time.time()
while frame_idx < args.max_frames:
log_probs = []
rewards = []
start_time = time.time()
for i_step in range(1, args.step_each_frame + 1):
# get initial attribute list
state = np.random.rand(1, args.num_inputs)
state = torch.from_numpy(state).float()
if cuda:
state = state.cuda()
# get modified attribute list
dist = main_model(state)
action = dist.sample()
action_actual = action.float() / 10.0 # [0, 0.9]
# generate images by attribute list
print("action: " + str(action_actual.cpu().numpy()))
get_images_by_attributes(args, i_step, env, brain,
action_actual[0].cpu().numpy())
train_set = get_training_dataset(args, i_step)
train_loader = data.DataLoader(train_set,
batch_size=args.syn_batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True)
# train the task model using synthetic dataset
task_model.load_state_dict(
torch.load(
os.path.join(args.snapshot_dir, "task_model_init.pth")))
reward = train_task_model(train_loader, val_loader, task_model,
task_optimizer, args, cuda)
log_prob = dist.log_prob(action)[0]
log_probs.append(log_prob)
rewards.append(torch.FloatTensor([reward]))
frame_idx += 1
if frame_idx == 1:
moving_start = torch.FloatTensor([reward])
baseline = compute_returns(rewards, moving_start)
moving_start = baseline[-1]
log_probs = torch.cat(log_probs)
baseline = torch.cat(baseline).detach()
rewards = torch.cat(rewards).detach()
advantage = rewards - baseline
if cuda:
advantage = advantage.cuda()
loss = -(log_probs * advantage.detach()).mean()
with open(os.path.join(args.snapshot_dir, "logs.txt"), 'a') as fp:
fp.write(
"frame idx: {0:4d}, state: {1:s}, action: {2:s}, reward: {3:s}, baseline: {4:s}, loss: {5:.2f} \n"
.format(frame_idx, str(state.cpu()[0].numpy()),
str(action.cpu()[0].numpy()), str(rewards.numpy()),
str(baseline.numpy()), loss.item()))
print("optimize the main model parameters")
main_optimizer.zero_grad()
loss.backward()
main_optimizer.step()
elapsed_time = time.time() - start_time
print("[frame: {0:3d}], [loss: {1:.2f}], [time: {2:.1f}]".format(
frame_idx, loss.item(), elapsed_time))
torch.save(
main_model.state_dict(),
os.path.join(args.snapshot_dir, "main_model_%d.pth" % frame_idx))
elapsed_time = time.time() - whole_start_time
print("whole time: {0:.1f}".format(elapsed_time))
env.close()
trainset = MNIST('./data', train=True, transform=trans_img)
testset = MNIST('./data', train=False, transform=trans_img)
trainloader = DataLoader(trainset,
batch_size=batch_size,
shuffle=True,
num_workers=4)
testloader = DataLoader(testset,
batch_size=batch_size,
shuffle=False,
num_workers=4)
# build network
mnist_mlp = MLP(io_bits=io_bits)
mnist_mlp.cuda()
try:
data = torch.load('mlp_mnist_nobias.t7')
mnist_mlp.load_state_dict(data)
print('Loading weight')
except:
print('Initializing model')
#
##
criterian = nn.CrossEntropyLoss(size_average=False)
optimizer = optim.SGD(mnist_mlp.parameters(), lr=learning_rate)
best_acc = 0
for i in range(epoches):
