Обратите внимание: 128, 129, 130 строк
Я использую memory_profiler, но у меня есть некоторые сомнения относительно ситуации с памятью, которую я наблюдал.Увеличение очень большое, но заполняемость не изменилась.Эта проблема вызвана проблемой memory_profiler или механизмом внутренней памяти python.
line Mem usage Increment Line Contents
================================================
58 675.719 MiB 675.719 MiB @profile
59 def train(dataset, embedding, tokenizer, entity_linker, min_count, max_word_length, max_entity_length, batch_size,
60 patience, optimizer, learning_rate, weight_decay, warmup_epochs, dropout_prob, use_gpu, use_word):
61 675.719 MiB 0.000 MiB batch_size = int(batch_size)
62 675.719 MiB 0.000 MiB learning_rate = float(learning_rate)
63 675.719 MiB 0.000 MiB weight_decay = float(weight_decay)
64 675.719 MiB 0.000 MiB if use_gpu:
65 device = torch.device('cuda')
66 else:
67 675.719 MiB 0.000 MiB device = torch.device('cpu')
68
69 947.828 MiB 272.109 MiB data = generate_features(dataset, tokenizer, entity_linker, min_count, max_word_length, max_entity_length)
70 947.828 MiB 0.000 MiB word_vocab = data['word_vocab']
71 947.828 MiB 0.000 MiB entity_vocab = data['entity_vocab']
72
73 947.828 MiB 0.000 MiB train_data_loader = DataLoader(data['train'], shuffle=True, batch_size=batch_size)
74 947.828 MiB 0.000 MiB dev_data_loader = DataLoader(data['dev'], shuffle=False, batch_size=batch_size)
75
76 947.828 MiB 0.000 MiB dim_size = embedding.syn0.shape[1]
77 976.418 MiB 28.590 MiB word_embedding = np.random.uniform(low=-0.05, high=0.05, size=(len(word_vocab), dim_size))
78 976.418 MiB 0.000 MiB word_embedding[0] = np.zeros(dim_size)
79 1043.434 MiB 0.000 MiB for word, index in word_vocab.items():
80 1043.434 MiB 0.000 MiB try:
81 1043.434 MiB 0.066 MiB word_embedding[index] = embedding.get_word_vector(word)
82 1043.434 MiB 0.000 MiB except KeyError:
83 1043.434 MiB 0.000 MiB continue
84 1099.840 MiB 56.406 MiB entity_embedding = np.random.uniform(low=-0.05, high=0.05, size=(len(entity_vocab), dim_size))
85 1099.840 MiB 0.000 MiB entity_embedding[0] = np.zeros(dim_size)
86 1213.281 MiB 0.000 MiB for entity, index in entity_vocab.items():
87 1213.281 MiB 0.000 MiB try:
88 1213.281 MiB 0.027 MiB entity_embedding[index] = embedding.get_entity_vector(entity)
89 1213.281 MiB 0.000 MiB except KeyError:
90 1213.281 MiB 0.000 MiB continue
91
92 1270.773 MiB 57.492 MiB model = NABoE(word_embedding, entity_embedding, len(dataset.label_names), dropout_prob, use_word)
93 1270.773 MiB 0.000 MiB if optimizer == 'adamW':
94 optimizer = AdamW(model.parameters(), lr=learning_rate, weight_decay=weight_decay,
95 warmup=warmup_epochs * len(train_data_loader))
96 1270.773 MiB 0.000 MiB elif optimizer == 'adam':
97 1270.773 MiB 0.000 MiB optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
98 elif optimizer == 'adagrad':
99 optimizer = torch.optim.Adagrad(model.parameters(), lr=learning_rate)
100 elif optimizer == 'sgd':
101 optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)
102 elif optimizer == 'radam':
103 optimizer = RAdam(model.parameters(), lr=learning_rate, weight_decay=weight_decay)
104 elif optimizer == 'ranger':
105 optimizer = Ranger(model.parameters(), lr=learning_rate, weight_decay=weight_decay)
106 1270.809 MiB 0.035 MiB model.to(device)
107
108 1270.809 MiB 0.000 MiB epoch = 0
109 1270.809 MiB 0.000 MiB best_val_acc = 0.0
110 1270.809 MiB 0.000 MiB best_weights = None
111 1270.809 MiB 0.000 MiB num_epochs_without_improvement = 0
112
113 1270.809 MiB 0.000 MiB total_memory = 0
114
115 1270.809 MiB 0.000 MiB while True:
116 1658.176 MiB 0.004 MiB with tqdm(train_data_loader) as pbar:
117 1658.176 MiB 0.000 MiB model.train()
118 1658.176 MiB 0.285 MiB for batch in pbar:
119
120 1658.176 MiB 0.000 MiB memory = cpu_stats()
121 # disk = psutil.disk_usage('/').percent
122 # cpu = psutil.cpu_percent(0)
123 # cpu_kernel = psutil.cpu_percent(percpu=True)
124
125 1658.176 MiB 0.000 MiB args = {k: v.to(device) for k, v in batch.items() if k != 'label'}
126 1892.559 MiB 263.688 MiB logits = model(**args)
127 1892.559 MiB 0.102 MiB loss = F.cross_entropy(logits, batch['label'].to(device))
128 1658.176 MiB 0.000 MiB loss.backward()
129 1658.176 MiB 127.637 MiB optimizer.step()
130 1658.176 MiB 0.008 MiB model.zero_grad()
131 1658.176 MiB 0.000 MiB total_memory += memory
132 1658.176 MiB 0.000 MiB pbar.set_description(f'epoch: {epoch} loss: {loss.item():.8f} The Memory Usage: {memory}'
133 f' Total Memory Usage: {total_memory}')
134
135 1658.176 MiB 0.000 MiB epoch += 1
136 1658.176 MiB 0.312 MiB val_acc = evaluate(model, dev_data_loader, device, 'dev')[0]
137 1658.176 MiB 0.000 MiB if val_acc > best_val_acc:
138 1658.176 MiB 0.000 MiB best_val_acc = val_acc
139 1658.176 MiB 0.004 MiB best_weights = {k: v.to('cpu').clone() for k, v in model.state_dict().items()}
140 1658.176 MiB 0.000 MiB num_epochs_without_improvement = 0
141 else:
142 1653.926 MiB 0.000 MiB num_epochs_without_improvement += 1
143
144 1658.176 MiB 0.000 MiB if num_epochs_without_improvement >= patience:
145 1653.859 MiB 0.000 MiB model.load_state_dict(best_weights)
146 1653.859 MiB 0.000 MiB break
147
148 1653.859 MiB 0.000 MiB test_data_loader = DataLoader(data['test'], shuffle=False, batch_size=batch_size) 149 1653.863 MiB 0.004 MiB return evaluate(model, test_data_loader, device, 'test')
129 строк увеличивают занятость на 127 МБ, но объем памяти не изменился, даже механизм внутренней переработки Python не совсем такой же