78 lines
2.5 KiB
Python
78 lines
2.5 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import imp
|
|
import numpy as np
|
|
|
|
class Base(nn.Module):
|
|
def __init__(self, stop_threshold):
|
|
super().__init__()
|
|
|
|
self.init_model()
|
|
self.num_params()
|
|
|
|
self.register_buffer("step", torch.zeros(1, dtype=torch.long))
|
|
self.register_buffer("stop_threshold", torch.tensor(stop_threshold, dtype=torch.float32))
|
|
|
|
@property
|
|
def r(self):
|
|
return self.decoder.r.item()
|
|
|
|
@r.setter
|
|
def r(self, value):
|
|
self.decoder.r = self.decoder.r.new_tensor(value, requires_grad=False)
|
|
|
|
def init_model(self):
|
|
for p in self.parameters():
|
|
if p.dim() > 1: nn.init.xavier_uniform_(p)
|
|
|
|
def finetune_partial(self, whitelist_layers):
|
|
self.zero_grad()
|
|
for name, child in self.named_children():
|
|
if name in whitelist_layers:
|
|
print("Trainable Layer: %s" % name)
|
|
print("Trainable Parameters: %.3f" % sum([np.prod(p.size()) for p in child.parameters()]))
|
|
for param in child.parameters():
|
|
param.requires_grad = False
|
|
|
|
def get_step(self):
|
|
return self.step.data.item()
|
|
|
|
def reset_step(self):
|
|
# assignment to parameters or buffers is overloaded, updates internal dict entry
|
|
self.step = self.step.data.new_tensor(1)
|
|
|
|
def log(self, path, msg):
|
|
with open(path, "a") as f:
|
|
print(msg, file=f)
|
|
|
|
def load(self, path, device, optimizer=None):
|
|
# Use device of model params as location for loaded state
|
|
checkpoint = torch.load(str(path), map_location=device)
|
|
if "model_state" in checkpoint:
|
|
state = checkpoint["model_state"]
|
|
else:
|
|
state = checkpoint["model"]
|
|
self.load_state_dict(state, strict=False)
|
|
|
|
if "optimizer_state" in checkpoint and optimizer is not None:
|
|
optimizer.load_state_dict(checkpoint["optimizer_state"])
|
|
|
|
def save(self, path, optimizer=None):
|
|
if optimizer is not None:
|
|
torch.save({
|
|
"model_state": self.state_dict(),
|
|
"optimizer_state": optimizer.state_dict(),
|
|
}, str(path))
|
|
else:
|
|
torch.save({
|
|
"model_state": self.state_dict(),
|
|
}, str(path))
|
|
|
|
|
|
def num_params(self, print_out=True):
|
|
parameters = filter(lambda p: p.requires_grad, self.parameters())
|
|
parameters = sum([np.prod(p.size()) for p in parameters]) / 1_000_000
|
|
if print_out:
|
|
print("Trainable Parameters: %.3fM" % parameters)
|
|
return parameters
|