model = Discriminator(args.length, len(species)).to(device)
optimizer = optim.Adam(model.parameters(), lr=args.rate)
# raise an error if receptive field is smaller than sampling length
if args.length < receptive_field(model):
raise Exception("Input sequences must be longer than {} bp.".format(
receptive_field(model)))
for epoch in range(args.epoch):
train(model, device, loader, optimizer, epoch + 1)
print("")
# calculate style matrices
if args.verbose > 1: print("Extracting style matrices...")
model.eval()
style_matrices = []
for record in SeqIO.parse(args.contig, "fasta"):
tensor = to_tensor(str(record.seq))
style_matrices += model.get_style(tensor.float().to(device),
args.layer)
style_matrices = torch.cat(style_matrices, dim=0)
torch.save(style_matrices, args.output)
if args.verbose > 1:
print("Genome style matrix is successfully written to {}.".format(
args.output))
