MCPcopy Create free account
hub / github.com/OpenRL-Lab/Wandb_Tutorial / test_pytorch

Function test_pytorch

basic/test_pytorch.py:37–95  ·  view source on GitHub ↗
(args)

Source from the content-addressed store, hash-verified

35 return all_args
36
37def test_pytorch(args):
38
39 all_args = parse(args)
40
41 # set all the seed
42 random.seed(all_args.seed)
43 torch.manual_seed(all_args.seed)
44 np.random.seed(all_args.seed)
45
46 run_dir = Path("../results") / all_args.project_name / all_args.experiment_name
47 if not run_dir.exists():
48 os.makedirs(str(run_dir))
49
50 wandb.init(config=all_args,
51 project=all_args.project_name,
52 entity=all_args.team_name,
53 notes=socket.gethostname(),
54 name=all_args.experiment_name + "_" + str(all_args.seed),
55 group=all_args.scenario_name,
56 dir=str(run_dir),
57 job_type="training",
58 reinit=True)
59
60 train_x = np.array([
61 [5.0, 3.5, 1.3, 0.3],
62 [4.5, 2.3, 1.3, 0.3],
63 [5.5, 2.6, 4.4, 1.2],
64 [6.1, 3.0, 4.6, 1.4],
65 [6.7, 3.1, 5.6, 2.4],
66 [6.9, 3.1, 5.1, 2.3]], dtype=np.float32)
67
68 train_y = np.array([0, 0, 1, 1, 2, 2], dtype=np.long)
69
70 train_x = torch.tensor(train_x, dtype=torch.float32)
71 train_y = torch.tensor(train_y, dtype=torch.long)
72 model = Model()
73 wandb.watch(model,log_freq=1)
74 max_epochs = 100
75 lrn_rate = 0.04
76 loss_func = torch.nn.CrossEntropyLoss()
77 optimizer = torch.optim.SGD(model.parameters(), lr=lrn_rate)
78
79 model.train()
80
81 indices = np.arange(6)
82 for epoch in range(0, max_epochs):
83 np.random.shuffle(indices)
84 for i in indices:
85 X = train_x[i].reshape(1, 4)
86 Y = train_y[i].reshape(1, )
87 optimizer.zero_grad()
88 oupt = model(X)
89 loss_obj = loss_func(oupt, Y)
90 loss_obj.backward()
91 optimizer.step()
92
93 print('Epoch:{} Loss:{}'.format(epoch,loss_obj))
94

Callers 1

test_pytorch.pyFile · 0.85

Calls 2

ModelClass · 0.85
parseFunction · 0.70

Tested by

no test coverage detected