Skip to content

Add support for initial configurations. #94

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
hzhangxyz wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Add support for initial configurations. #94

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 14 additions & 0 deletions qmb/haar.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
from .subcommand_dict import subcommand_dict
from .model_dict import ModelProto
from .optimizer import initialize_optimizer, scale_learning_rate
from .bitspack import pack_int


@dataclasses.dataclass
Expand Down Expand Up @@ -256,6 +257,8 @@ class HaarConfig:

common: typing.Annotated[CommonConfig, tyro.conf.OmitArgPrefixes]

# The initial configurations for the first step
initial_config: typing.Annotated[str, tyro.conf.arg(aliases=["-i"])] = ""
# The sampling count from neural network
sampling_count_from_neural_network: typing.Annotated[int, tyro.conf.arg(aliases=["-n"])] = 1024
# The sampling count from last iteration
Expand Down Expand Up @@ -364,6 +367,17 @@ def main(self, *, model_param: typing.Any = None, network_param: typing.Any = No
pool_configs, pool_psi = data["haar"]["pool"]
data["haar"]["pool"] = (pool_configs.to(device=self.common.device), pool_psi.to(device=self.common.device))

if self.initial_config != "":
if data["haar"]["pool"] is None:
config = pack_int(
torch.tensor([[int(i) for i in single_config] for single_config in self.initial_config.split(",")], dtype=torch.uint8, device=self.common.device),
size=1,
)
data["haar"]["pool"] = (config, network(config))
logging.info("The initial configuration is imported successfully.")
else:
logging.info("The initial configuration is provided, but the pool from the last iteration is not empty, so the initial configuration will be ignored.")

writer = torch.utils.tensorboard.SummaryWriter(log_dir=self.common.folder()) # type: ignore[no-untyped-call]

while True:
Expand Down