[WandB] wandb artifacts를 vscode에서 사용하는 방법

qw4735·2023년 3월 20일
0

WandB

목록 보기
2/3

wandb로 돌린 모델 사용하는 방법


예를 들어, seed가 9인 version 85를 돌리고 싶으면, 해당 Artifacts에 Usage에 들어가서 해당 코드를 써준다.

import wandb
run = wandb.init()
artifact = run.use_artifact('qw4735/TVAE/TVAE_Personal_Loan:v85', type='model')
artifact_dir = artifact.download()
######
set_random_seed(config['seed'])
torch.manual_seed(config['seed'])
if config['cuda']:
    torch.cuda.manual_seed(config['seed'])

model = TVAE(config, device).to(device)

if config["cuda"]:
    model_name = [x for x in os.listdir(artifact_dir) if x.endswith('pth')][0]
    model.load_state_dict(torch.load(artifact_dir + '/' + model_name))
else:
    model_name = [x for x in os.listdir(artifact_dir) if x.endswith('pth')][0]
    model.load_state_dict(torch.load(artifact_dir + '/' + model_name, map_location=torch.device('cpu')))

model.eval()

0개의 댓글

관련 채용 정보