from pathlib import Path
from deep_t2i.trainer_DAMSM import BirdsTrainer

data_dir = Path('/root/data/birds')
model_dir = Path('../models/large_birds/damsm')
data_dir, model_dir

(Path('/root/data/birds'), Path('../models/large_birds/damsm'))

Train

bs = 48
data_pct = 1
lr = 2e-3
lr_decay = 0.98
emb_sz = 256
rnn_layers = 2
rnn_drop_p = 0.5
gamma1 = 4.0
gamma2 = 5.0
gamma3 = 10.0
trainer = BirdsTrainer(
    data_dir, 
    bs=bs, 
    data_pct=data_pct,
    lr=lr, 
    lr_decay=0.98,
    device='cuda', 
    emb_sz=emb_sz,
    rnn_layers=rnn_layers,
    rnn_drop_p=rnn_drop_p,
    gamma1=gamma1,
    gamma2=gamma2,
    gamma3=gamma3,
)

step_per_epoch = len(trainer.dls.train)
step_per_epoch

245

trainer.train(
    step_per_epoch*230, 
    step_per_epoch=step_per_epoch,
    saveck_every=step_per_epoch*10, 
    ck_path=str(model_dir/'0'), 
)

1, time: 94.2s, loss: 13.8903
2, time: 89.8s, loss: 11.3273
3, time: 89.7s, loss: 10.4841
4, time: 89.9s, loss: 9.9330
5, time: 90.3s, loss: 9.5434
6, time: 89.2s, loss: 9.3482
7, time: 89.3s, loss: 9.0983
8, time: 89.5s, loss: 8.8963
9, time: 89.5s, loss: 8.7451
10, time: 89.9s, loss: 8.6581
11, time: 90.9s, loss: 8.5409
12, time: 91.6s, loss: 8.5167
13, time: 89.7s, loss: 8.3063
14, time: 89.4s, loss: 8.3121
15, time: 89.4s, loss: 8.1490
16, time: 89.4s, loss: 8.1143
17, time: 90.0s, loss: 8.0694
18, time: 90.6s, loss: 8.0042
19, time: 89.8s, loss: 7.9207
20, time: 90.5s, loss: 7.9125
21, time: 89.6s, loss: 7.8875
22, time: 89.0s, loss: 7.7886
23, time: 88.6s, loss: 7.6879
24, time: 88.5s, loss: 7.7285
25, time: 88.4s, loss: 7.6062
26, time: 88.1s, loss: 7.6070
27, time: 88.0s, loss: 7.5772
28, time: 88.2s, loss: 7.5785
29, time: 88.3s, loss: 7.5019
30, time: 88.5s, loss: 7.4777
31, time: 88.3s, loss: 7.3824
32, time: 88.1s, loss: 7.3545
33, time: 87.9s, loss: 7.2782
34, time: 88.2s, loss: 7.2882
35, time: 87.7s, loss: 7.2488
36, time: 87.8s, loss: 7.2347
37, time: 88.2s, loss: 7.2094
38, time: 87.4s, loss: 7.1717
39, time: 87.7s, loss: 7.1521
40, time: 87.5s, loss: 7.1742
41, time: 87.5s, loss: 7.0289
42, time: 87.2s, loss: 7.0144
43, time: 87.2s, loss: 7.0822
44, time: 87.5s, loss: 7.0336
45, time: 87.6s, loss: 6.9589
46, time: 87.5s, loss: 6.9306
47, time: 87.8s, loss: 6.9804
48, time: 89.4s, loss: 6.9057
49, time: 90.2s, loss: 6.8482
50, time: 90.1s, loss: 6.8292
51, time: 88.6s, loss: 6.8695
52, time: 93.7s, loss: 6.7394
53, time: 90.8s, loss: 6.7923
54, time: 89.4s, loss: 6.7830
55, time: 88.7s, loss: 6.7368
56, time: 88.9s, loss: 6.7868
57, time: 88.9s, loss: 6.7407
58, time: 88.8s, loss: 6.6657
59, time: 89.1s, loss: 6.6805
60, time: 89.1s, loss: 6.6612
61, time: 89.0s, loss: 6.6024
62, time: 88.8s, loss: 6.6956
63, time: 88.7s, loss: 6.6074
64, time: 89.0s, loss: 6.6055
65, time: 89.2s, loss: 6.6176
66, time: 88.5s, loss: 6.5470
67, time: 88.2s, loss: 6.6552
68, time: 87.9s, loss: 6.6009
69, time: 90.3s, loss: 6.5428
70, time: 93.9s, loss: 6.5057
71, time: 93.6s, loss: 6.4873
72, time: 91.7s, loss: 6.5172
73, time: 89.8s, loss: 6.4443
74, time: 92.2s, loss: 6.3770
75, time: 89.5s, loss: 6.4399
76, time: 89.0s, loss: 6.4024
77, time: 88.2s, loss: 6.4142
78, time: 88.3s, loss: 6.4230
79, time: 87.4s, loss: 6.4008
80, time: 87.8s, loss: 6.4197
81, time: 88.3s, loss: 6.3578
82, time: 88.4s, loss: 6.3773
83, time: 87.9s, loss: 6.2574
84, time: 87.4s, loss: 6.2924
85, time: 88.1s, loss: 6.3449
86, time: 87.5s, loss: 6.2831
87, time: 87.6s, loss: 6.3095
88, time: 88.0s, loss: 6.2992
89, time: 87.4s, loss: 6.3533
90, time: 93.0s, loss: 6.2908
91, time: 88.1s, loss: 6.2466
92, time: 87.7s, loss: 6.1447
93, time: 87.5s, loss: 6.2131
94, time: 87.5s, loss: 6.2334
95, time: 87.9s, loss: 6.2098
96, time: 87.6s, loss: 6.2295
97, time: 87.4s, loss: 6.1301
98, time: 87.4s, loss: 6.1103
99, time: 87.5s, loss: 6.1902
100, time: 87.8s, loss: 6.2157
101, time: 87.9s, loss: 6.1310
102, time: 87.9s, loss: 6.1695
103, time: 87.5s, loss: 6.1448
104, time: 87.4s, loss: 6.1302
105, time: 87.7s, loss: 6.0903
106, time: 87.9s, loss: 6.0956
107, time: 87.7s, loss: 6.1523
108, time: 87.9s, loss: 6.0722
109, time: 87.9s, loss: 6.0711
110, time: 87.7s, loss: 6.0791
111, time: 87.5s, loss: 6.0297
112, time: 87.3s, loss: 6.1128
113, time: 87.4s, loss: 6.0748
114, time: 87.1s, loss: 6.0337
115, time: 87.9s, loss: 6.1036
116, time: 88.1s, loss: 6.1205
117, time: 87.9s, loss: 6.0527
118, time: 89.0s, loss: 6.0014
119, time: 90.3s, loss: 6.0399
120, time: 92.8s, loss: 6.0535
121, time: 92.7s, loss: 6.0148
122, time: 92.5s, loss: 6.0990
123, time: 92.8s, loss: 6.0407
124, time: 93.4s, loss: 6.0545
125, time: 93.3s, loss: 6.0111
126, time: 93.6s, loss: 6.0261
127, time: 93.8s, loss: 5.9871
128, time: 93.6s, loss: 6.0813
129, time: 94.2s, loss: 5.9546
130, time: 94.5s, loss: 5.9433
131, time: 93.5s, loss: 6.0418
132, time: 94.0s, loss: 5.9356
133, time: 94.4s, loss: 5.8967
134, time: 93.8s, loss: 5.9490
135, time: 94.5s, loss: 6.0175
136, time: 94.3s, loss: 5.9631
137, time: 94.1s, loss: 5.9302
138, time: 94.7s, loss: 6.0012
139, time: 94.2s, loss: 5.9189
140, time: 94.3s, loss: 5.9537
141, time: 94.2s, loss: 5.9725
142, time: 94.2s, loss: 5.9390
143, time: 94.1s, loss: 5.9698
144, time: 93.8s, loss: 6.0037
145, time: 93.8s, loss: 5.8987
146, time: 94.0s, loss: 5.9426
147, time: 94.0s, loss: 5.9372
148, time: 94.1s, loss: 5.9323
149, time: 94.3s, loss: 5.9987
150, time: 94.8s, loss: 5.8736
151, time: 94.2s, loss: 5.9392
152, time: 94.3s, loss: 5.9153
153, time: 94.6s, loss: 5.9035
154, time: 94.1s, loss: 5.9264
155, time: 93.9s, loss: 5.8856
156, time: 93.9s, loss: 5.8761
157, time: 93.8s, loss: 5.9274
158, time: 93.9s, loss: 5.8041
159, time: 94.0s, loss: 5.8952
160, time: 95.0s, loss: 5.9461
161, time: 94.8s, loss: 5.7963
162, time: 94.4s, loss: 5.8743
163, time: 94.6s, loss: 5.8709
164, time: 94.6s, loss: 5.8664
165, time: 94.5s, loss: 5.8318
166, time: 95.0s, loss: 5.7663
167, time: 94.4s, loss: 5.7967
168, time: 94.3s, loss: 5.8426
169, time: 94.3s, loss: 5.8984
170, time: 95.0s, loss: 5.8181
171, time: 94.8s, loss: 5.8536
172, time: 94.5s, loss: 5.8188
173, time: 94.6s, loss: 5.7263
174, time: 94.3s, loss: 5.8249
175, time: 94.1s, loss: 5.7796
176, time: 94.4s, loss: 5.7599
177, time: 94.3s, loss: 5.8192
178, time: 94.3s, loss: 5.8556
179, time: 94.7s, loss: 5.8563
180, time: 95.1s, loss: 5.7447
181, time: 95.1s, loss: 5.7788
182, time: 94.2s, loss: 5.8042
183, time: 94.2s, loss: 5.7773
184, time: 94.4s, loss: 5.8318
185, time: 94.7s, loss: 5.7647
186, time: 94.0s, loss: 5.7676
187, time: 94.3s, loss: 5.7583
188, time: 94.4s, loss: 5.7572
189, time: 94.5s, loss: 5.6973
190, time: 94.9s, loss: 5.8566
191, time: 94.8s, loss: 5.7318
192, time: 94.0s, loss: 5.7484
193, time: 94.2s, loss: 5.7182
194, time: 94.6s, loss: 5.7778
195, time: 94.5s, loss: 5.7247
196, time: 94.7s, loss: 5.7313
197, time: 94.5s, loss: 5.7468
198, time: 94.2s, loss: 5.7163
199, time: 94.7s, loss: 5.6732
200, time: 95.0s, loss: 5.7757
201, time: 95.4s, loss: 5.6893
202, time: 94.4s, loss: 5.7395
203, time: 94.7s, loss: 5.7362
204, time: 94.7s, loss: 5.7453
205, time: 94.6s, loss: 5.7289
206, time: 94.2s, loss: 5.6784
207, time: 94.3s, loss: 5.7390
208, time: 94.0s, loss: 5.6593
209, time: 94.5s, loss: 5.6708
210, time: 94.8s, loss: 5.6992
211, time: 94.9s, loss: 5.6877
212, time: 94.9s, loss: 5.6946
213, time: 94.5s, loss: 5.7151
214, time: 94.7s, loss: 5.6562
215, time: 94.3s, loss: 5.7091
216, time: 94.4s, loss: 5.6948
217, time: 94.7s, loss: 5.6774
218, time: 94.7s, loss: 5.6890
219, time: 94.7s, loss: 5.6532
220, time: 95.0s, loss: 5.6362
221, time: 95.0s, loss: 5.5935
222, time: 94.6s, loss: 5.7008
223, time: 94.3s, loss: 5.6589
224, time: 95.1s, loss: 5.6505
225, time: 95.3s, loss: 5.5928
226, time: 94.6s, loss: 5.6684
227, time: 94.7s, loss: 5.6739
228, time: 95.1s, loss: 5.6370
229, time: 94.4s, loss: 5.6772
230, time: 94.6s, loss: 5.6463

total_time: 350.7min

trainer.load_checkpoint(model_dir/'0-23.pt')

trainer.show()

trainer.show()

trainer.export(model_dir/'damsm_export.pt')

# from deep_t2i.model import get_pretrained_DAMSM
# rnn_encoder, cnn_encoder, gamma1, gamma2, gamma3 = get_pretrained_DAMSM(model_dir/'damsm_export.pt', device='cuda')

# trainer.rnn_encoder = rnn_encoder
# trainer.cnn_encoder = cnn_encoder